Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(jax/array-api): DPA-2 #4294

Merged
merged 7 commits into from
Nov 2, 2024
Merged
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 25 additions & 12 deletions deepmd/dpmodel/descriptor/dpa2.py
Original file line number Diff line number Diff line change
@@ -4,11 +4,18 @@
Union,
)

import array_api_compat
import numpy as np

from deepmd.dpmodel import (
NativeOP,
)
from deepmd.dpmodel.array_api import (
xp_take_along_axis,
)
from deepmd.dpmodel.common import (
to_numpy_array,
)
from deepmd.dpmodel.utils import (
EnvMat,
NetworkCollection,
@@ -787,9 +794,10 @@ def call(
The smooth switch function. shape: nf x nloc x nnei

"""
xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
use_three_body = self.use_three_body
nframes, nloc, nnei = nlist.shape
nall = coord_ext.reshape(nframes, -1).shape[1] // 3
nall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3
# nlists
nlist_dict = build_multiple_neighbor_list(
coord_ext,
@@ -798,7 +806,10 @@ def call(
self.nsel_list,
)
# repinit
g1_ext = self.type_embedding.call()[atype_ext]
g1_ext = xp.reshape(
xp.take(self.type_embedding.call(), xp.reshape(atype_ext, [-1]), axis=0),
(nframes, nall, self.tebd_dim),
)
g1_inp = g1_ext[:, :nloc, :]
g1, _, _, _, _ = self.repinit(
nlist_dict[
@@ -823,16 +834,18 @@ def call(
g1_ext,
mapping,
)
g1 = np.concatenate([g1, g1_three_body], axis=-1)
g1 = xp.concat([g1, g1_three_body], axis=-1)
# linear to change shape
g1 = self.g1_shape_tranform(g1)
if self.add_tebd_to_repinit_out:
assert self.tebd_transform is not None
g1 = g1 + self.tebd_transform(g1_inp)
# mapping g1
assert mapping is not None
mapping_ext = np.tile(mapping.reshape(nframes, nall, 1), (1, 1, g1.shape[-1]))
g1_ext = np.take_along_axis(g1, mapping_ext, axis=1)
mapping_ext = xp.tile(
xp.reshape(mapping, (nframes, nall, 1)), (1, 1, g1.shape[-1])
)
g1_ext = xp_take_along_axis(g1, mapping_ext, axis=1)
# repformer
g1, g2, h2, rot_mat, sw = self.repformers(
nlist_dict[
@@ -846,7 +859,7 @@ def call(
mapping,
)
if self.concat_output_tebd:
g1 = np.concatenate([g1, g1_inp], axis=-1)
g1 = xp.concat([g1, g1_inp], axis=-1)
return g1, rot_mat, g2, h2, sw

def serialize(self) -> dict:
@@ -883,8 +896,8 @@ def serialize(self) -> dict:
"embeddings": repinit.embeddings.serialize(),
"env_mat": EnvMat(repinit.rcut, repinit.rcut_smth).serialize(),
"@variables": {
"davg": repinit["davg"],
"dstd": repinit["dstd"],
"davg": to_numpy_array(repinit["davg"]),
"dstd": to_numpy_array(repinit["dstd"]),
},
}
if repinit.tebd_input_mode in ["strip"]:
@@ -896,8 +909,8 @@ def serialize(self) -> dict:
"repformer_layers": [layer.serialize() for layer in repformers.layers],
"env_mat": EnvMat(repformers.rcut, repformers.rcut_smth).serialize(),
"@variables": {
"davg": repformers["davg"],
"dstd": repformers["dstd"],
"davg": to_numpy_array(repformers["davg"]),
"dstd": to_numpy_array(repformers["dstd"]),
},
}
data.update(
@@ -913,8 +926,8 @@ def serialize(self) -> dict:
repinit_three_body.rcut, repinit_three_body.rcut_smth
).serialize(),
"@variables": {
"davg": repinit_three_body["davg"],
"dstd": repinit_three_body["dstd"],
"davg": to_numpy_array(repinit_three_body["davg"]),
"dstd": to_numpy_array(repinit_three_body["dstd"]),
},
}
if repinit_three_body.tebd_input_mode in ["strip"]:
Loading