feat(c++,pt-expt): add .pt2 (AOTInductor) C/C++ inference with DPA1/DPA2/DPA3 support

deepmd/backend/pt_expt.py

@@ -41,7 +41,7 @@ class PyTorchExportableBackend(Backend):
         | Backend.Feature.IO
     )
     """The features of the backend."""
-    suffixes: ClassVar[list[str]] = [".pte"]
+    suffixes: ClassVar[list[str]] = [".pte", ".pt2"]
     """The suffixes of the backend."""
 
     def is_available(self) -> bool:

deepmd/dpmodel/descriptor/dpa1.py

@@ -20,6 +20,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     cast_precision,
@@ -536,7 +537,7 @@ def call(
             (nf, nall, self.tebd_dim),
         )
         # nfnl x tebd_dim
-        atype_embd = atype_embd_ext[:, :nloc, :]
+        atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc)
         grrg, g2, h2, rot_mat, sw = self.se_atten(
             nlist,
             coord_ext,
@@ -1086,7 +1087,7 @@ def call(
             self.stddev[...],
         )
         nf, nloc, nnei, _ = dmatrix.shape
-        atype = atype_ext[:, :nloc]
+        atype = xp_take_first_n(atype_ext, 1, nloc)
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # nfnl x nnei
         exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
@@ -1105,6 +1106,12 @@ def call(
         nlist_masked = xp.where(nlist_mask, nlist, xp.zeros_like(nlist))
         ng = self.neuron[-1]
         nt = self.tebd_dim
+
+        # Gather neighbor info using xp_take_along_axis along axis=1.
+        # This avoids flat (nf*nall,) indexing that creates Ne(nall, nloc)
+        # constraints in torch.export, breaking NoPbc (nall == nloc).
+        nlist_2d = xp.reshape(nlist_masked, (nf, nloc * nnei))  # (nf, nloc*nnei)
+
         # nfnl x nnei x 4
         rr = xp.reshape(dmatrix, (nf * nloc, nnei, 4))
         rr = rr * xp.astype(exclude_mask[:, :, None], rr.dtype)
@@ -1113,15 +1120,16 @@ def call(
         if self.tebd_input_mode in ["concat"]:
             # nfnl x tebd_dim
             atype_embd = xp.reshape(
-                atype_embd_ext[:, :nloc, :], (nf * nloc, self.tebd_dim)
+                xp_take_first_n(atype_embd_ext, 1, nloc), (nf * nloc, self.tebd_dim)
             )
             # nfnl x nnei x tebd_dim
             atype_embd_nnei = xp.tile(atype_embd[:, xp.newaxis, :], (1, nnei, 1))
-            index = xp.tile(
-                xp.reshape(nlist_masked, (nf, -1, 1)), (1, 1, self.tebd_dim)
+            # Gather neighbor type embeddings: (nf, nall, tebd_dim) -> (nf, nloc*nnei, tebd_dim)
+            nlist_idx_tebd = xp.tile(nlist_2d[:, :, xp.newaxis], (1, 1, self.tebd_dim))
+            atype_embd_nlist = xp_take_along_axis(
+                atype_embd_ext, nlist_idx_tebd, axis=1
             )
             # nfnl x nnei x tebd_dim
-            atype_embd_nlist = xp_take_along_axis(atype_embd_ext, index, axis=1)
             atype_embd_nlist = xp.reshape(
                 atype_embd_nlist, (nf * nloc, nnei, self.tebd_dim)
             )
@@ -1140,10 +1148,9 @@ def call(
             assert self.embeddings_strip is not None
             assert type_embedding is not None
             ntypes_with_padding = type_embedding.shape[0]
-            # nf x (nl x nnei)
-            nlist_index = xp.reshape(nlist_masked, (nf, nloc * nnei))
-            # nf x (nl x nnei)
-            nei_type = xp_take_along_axis(atype_ext, nlist_index, axis=1)
+            # Gather neighbor types: (nf, nall) -> (nf, nloc*nnei)
+            nei_type = xp_take_along_axis(atype_ext, nlist_2d, axis=1)
+            nei_type = xp.reshape(nei_type, (-1,))  # (nf * nloc * nnei,)
             # (nf x nl x nnei) x ng
             nei_type_index = xp.tile(xp.reshape(nei_type, (-1, 1)), (1, ng))
             if self.type_one_side:

deepmd/dpmodel/descriptor/dpa2.py

@@ -15,6 +15,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     cast_precision,
@@ -878,7 +879,7 @@ def call(
             xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
             (nframes, nall, self.tebd_dim),
         )
-        g1_inp = g1_ext[:, :nloc, :]
+        g1_inp = xp_take_first_n(g1_ext, 1, nloc)
         g1, _, _, _, _ = self.repinit(
             nlist_dict[
                 get_multiple_nlist_key(self.repinit.get_rcut(), self.repinit.get_nsel())
@@ -912,9 +913,7 @@ def call(
             g1 = g1 + self.tebd_transform(g1_inp)
         # mapping g1
         assert mapping is not None
-        mapping_ext = xp.tile(
-            xp.reshape(mapping, (nframes, nall, 1)), (1, 1, g1.shape[-1])
-        )
+        mapping_ext = xp.tile(xp.expand_dims(mapping, axis=-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(

deepmd/dpmodel/descriptor/dpa3.py

@@ -10,6 +10,7 @@
 )
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     cast_precision,
@@ -653,7 +654,11 @@ def call(
         type_embedding = self.type_embedding.call()
         if self.use_loc_mapping:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0),
+                xp.take(
+                    type_embedding,
+                    xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (-1,)),
+                    axis=0,
+                ),
                 (nframes, nloc, self.tebd_dim),
             )
         else:
@@ -682,7 +687,7 @@ def call(
             sys_cs_embd = self.cs_activation_fn(self.mix_cs_mlp.call(cs_cat))
             node_ebd_ext = node_ebd_ext + xp.expand_dims(sys_cs_embd, axis=1)
 
-        node_ebd_inp = node_ebd_ext[:, :nloc, :]
+        node_ebd_inp = xp_take_first_n(node_ebd_ext, 1, nloc)
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(
             nlist,

deepmd/dpmodel/descriptor/repflows.py

@@ -13,6 +13,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -562,7 +563,7 @@ def call(
 
         # get node embedding
         # nb x nloc x tebd_dim
-        atype_embd = atype_embd_ext[:, :nloc, :]
+        atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc)
         assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
 
         node_ebd = self.act(atype_embd)
@@ -641,7 +642,7 @@ def call(
         angle_ebd = self.angle_embd(angle_input)
 
         # nb x nall x n_dim
-        mapping = xp.tile(xp.reshape(mapping, (nframes, -1, 1)), (1, 1, self.n_dim))
+        mapping = xp.tile(xp.expand_dims(mapping, axis=-1), (1, 1, self.n_dim))
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim
@@ -1421,7 +1422,7 @@ def call(
         n_edge = (
             int(xp.sum(xp.astype(nlist_mask, xp.int32))) if self.use_dynamic_sel else 0
         )
-        node_ebd = node_ebd_ext[:, :nloc, :]
+        node_ebd = xp_take_first_n(node_ebd_ext, 1, nloc)
         assert (nb, nloc) == node_ebd.shape[:2]
         if not self.use_dynamic_sel:
             assert (nb, nloc, nnei) == h2.shape[:3]

deepmd/dpmodel/descriptor/repformers.py

@@ -16,6 +16,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -499,7 +500,7 @@ def call(
         sw = xp.reshape(sw, (nf, nloc, nnei))
         sw = xp.where(nlist_mask, sw, xp.zeros_like(sw))
         # nf x nloc x tebd_dim
-        atype_embd = atype_embd_ext[:, :nloc, :]
+        atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc)
         assert list(atype_embd.shape) == [nf, nloc, self.g1_dim]
 
         g1 = self.act(atype_embd)
@@ -516,7 +517,7 @@ def call(
         # if a neighbor is real or not is indicated by nlist_mask
         nlist = xp.where(nlist == -1, xp.zeros_like(nlist), nlist)
         # nf x nall x ng1
-        mapping = xp.tile(xp.reshape(mapping, (nf, -1, 1)), (1, 1, self.g1_dim))
+        mapping = xp.tile(xp.expand_dims(mapping, axis=-1), (1, 1, self.g1_dim))
         for idx, ll in enumerate(self.layers):
             # g1:     nf x nloc x ng1
             # g1_ext: nf x nall x ng1
@@ -1765,9 +1766,8 @@ def call(
         )
 
         nf, nloc, nnei, _ = g2.shape
-        nall = g1_ext.shape[1]
         # g1, _ = xp.split(g1_ext, [nloc], axis=1)
-        g1 = g1_ext[:, :nloc, :]
+        g1 = xp_take_first_n(g1_ext, 1, nloc)
         assert (nf, nloc) == g1.shape[:2]
         assert (nf, nloc, nnei) == h2.shape[:3]
 

deepmd/dpmodel/model/make_model.py

@@ -11,6 +11,8 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.atomic_model.base_atomic_model import (
     BaseAtomicModel,
@@ -589,7 +591,6 @@ def _format_nlist(
             xp = array_api_compat.array_namespace(extended_coord, nlist)
             n_nf, n_nloc, n_nnei = nlist.shape
             extended_coord = extended_coord.reshape([n_nf, -1, 3])
-            nall = extended_coord.shape[1]
             rcut = self.get_rcut()
 
             if n_nnei < nnei:
@@ -612,14 +613,14 @@ def _format_nlist(
                 # make a copy before revise
                 m_real_nei = nlist >= 0
                 ret = xp.where(m_real_nei, nlist, 0)
-                coord0 = extended_coord[:, :n_nloc, :]
+                coord0 = xp_take_first_n(extended_coord, 1, n_nloc)
                 index = xp.tile(ret.reshape(n_nf, n_nloc * n_nnei, 1), (1, 1, 3))
-                coord1 = xp.take_along_axis(extended_coord, index, axis=1)
+                coord1 = xp_take_along_axis(extended_coord, index, axis=1)
                 coord1 = coord1.reshape(n_nf, n_nloc, n_nnei, 3)
                 rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1)
                 rr = xp.where(m_real_nei, rr, float("inf"))
                 rr, ret_mapping = xp.sort(rr, axis=-1), xp.argsort(rr, axis=-1)
-                ret = xp.take_along_axis(ret, ret_mapping, axis=2)
+                ret = xp_take_along_axis(ret, ret_mapping, axis=2)
                 ret = xp.where(rr > rcut, -1, ret)
                 ret = ret[..., :nnei]
             # not extra_nlist_sort and n_nnei <= nnei:

deepmd/dpmodel/utils/exclude_mask.py

@@ -6,6 +6,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 
 
@@ -131,18 +132,22 @@ def build_type_exclude_mask(
             ],
             axis=-1,
         )
-        type_i = xp.reshape(atype_ext[:, :nloc], (nf, nloc)) * (self.ntypes + 1)
-        # nf x nloc x nnei
-        index = xp.reshape(
-            xp.where(nlist == -1, xp.full_like(nlist, nall), nlist), (nf, nloc * nnei)
+        type_i = xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (nf, nloc)) * (
+            self.ntypes + 1
         )
-        type_j = xp_take_along_axis(ae, index, axis=1)
+        # Map -1 entries to nall (the virtual atom index in ae)
+        nlist_for_type = xp.where(nlist == -1, xp.full_like(nlist, nall), nlist)
+        # Gather neighbor types using xp_take_along_axis along axis=1.
+        # This avoids flat (nf*(nall+1),) indexing that creates Ne(nall, nloc)
+        # constraints in torch.export, breaking NoPbc (nall == nloc).
+        nlist_for_gather = xp.reshape(nlist_for_type, (nf, nloc * nnei))
+        type_j = xp_take_along_axis(ae, nlist_for_gather, axis=1)
         type_j = xp.reshape(type_j, (nf, nloc, nnei))
         type_ij = type_i[:, :, None] + type_j
-        # nf x (nloc x nnei)
-        type_ij = xp.reshape(type_ij, (nf, nloc * nnei))
+        # (nf * nloc * nnei,)
+        type_ij_flat = xp.reshape(type_ij, (-1,))
         mask = xp.reshape(
-            xp.take(self.type_mask[...], xp.reshape(type_ij, (-1,))),
+            xp.take(self.type_mask[...], type_ij_flat),
             (nf, nloc, nnei),
         )
         return mask

deepmd/dpmodel/utils/nlist.py

@@ -5,6 +5,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 
 from .region import (
@@ -243,8 +244,7 @@ def build_multiple_neighbor_list(
         nlist = xp.concat([nlist, pad], axis=-1)
         nsel = nsels[-1]
     coord1 = xp.reshape(coord, (nb, -1, 3))
-    nall = coord1.shape[1]
-    coord0 = coord1[:, :nloc, :]
+    coord0 = xp_take_first_n(coord1, 1, nloc)
     nlist_mask = nlist == -1
     tnlist_0 = xp.where(nlist_mask, xp.zeros_like(nlist), nlist)
     index = xp.tile(xp.reshape(tnlist_0, (nb, nloc * nsel, 1)), (1, 1, 3))

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -73,6 +73,23 @@
     extend_descrpt_stat,
 )
 
+if not hasattr(torch.ops.deepmd, "tabulate_fusion_se_t_tebd"):
+
+    def tabulate_fusion_se_t_tebd(
+        argument0: torch.Tensor,
+        argument1: torch.Tensor,
+        argument2: torch.Tensor,
+        argument3: torch.Tensor,
+        argument4: int,
+    ) -> list[torch.Tensor]:
+        raise NotImplementedError(
+            "tabulate_fusion_se_t_tebd is not available since customized PyTorch OP library is not built when freezing the model. "
+            "See documentation for model compression for details."
+        )
+
+    # Note: this hack cannot actually save a model that can be run using LAMMPS.
+    torch.ops.deepmd.tabulate_fusion_se_t_tebd = tabulate_fusion_se_t_tebd
+
 
 @BaseDescriptor.register("se_e3_tebd")
 class DescrptSeTTebd(BaseDescriptor, torch.nn.Module):