feat(pt_expt): add dp freeze support for pt_expt backend (.pte/.pt2)

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,
@@ -534,7 +535,7 @@
             (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,
@@ -1056,7 +1057,8 @@
             self.stddev[...],
         )
         nf, nloc, nnei, _ = dmatrix.shape
-        atype = atype_ext[:, :nloc]
+        nall = atype_ext.shape[1]
+        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))
@@ -1075,6 +1077,12 @@
         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)
@@ -1083,15 +1091,16 @@
         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)
             )
@@ -1110,10 +1119,9 @@
             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,
@@ -876,7 +877,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())
@@ -910,9 +911,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/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,
@@ -558,7 +560,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:
@@ -581,14 +582,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_expt/entrypoints/main.py

@@ -160,6 +160,54 @@ def train(
     trainer.run()
 
 
+def freeze(
+    model: str,
+    output: str = "frozen_model.pt2",
+    head: str | None = None,
+) -> None:
+    """Freeze a pt_expt training checkpoint to .pte or .pt2 format.
+
+    Parameters
+    ----------
+    model : str
+        Path to the training checkpoint (.pt file).
+    output : str
+        Path for the frozen model output (.pte or .pt2).
+    head : str or None
+        Head to freeze in a multi-task model (not yet supported).
+    """
+    import torch
+
+    from deepmd.pt_expt.model import (
+        get_model,
+    )
+    from deepmd.pt_expt.train.wrapper import (
+        ModelWrapper,
+    )
+    from deepmd.pt_expt.utils.env import (
+        DEVICE,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+    )
+
+    state_dict = torch.load(model, map_location=DEVICE, weights_only=True)
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+    model_params = state_dict["_extra_state"]["model_params"]
+
+    # Reconstruct model and load weights
+    pt_expt_model = get_model(model_params).to(DEVICE)
+    wrapper = ModelWrapper(pt_expt_model)
+    wrapper.load_state_dict(state_dict)
+    pt_expt_model.eval()
+
+    # Serialize to dict and export
+    model_dict = pt_expt_model.serialize()
+    deserialize_to_file(output, {"model": model_dict})
+    log.info(f"Saved frozen model to {output}")
+
+
 def main(args: list[str] | argparse.Namespace | None = None) -> None:
     """Entry point for the pt_expt backend CLI.
 
@@ -195,6 +243,18 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None:
             skip_neighbor_stat=FLAGS.skip_neighbor_stat,
             output=FLAGS.output,
         )
+    elif FLAGS.command == "freeze":
+        if Path(FLAGS.checkpoint_folder).is_dir():
+            checkpoint_path = Path(FLAGS.checkpoint_folder)
+            latest_ckpt_file = (checkpoint_path / "checkpoint").read_text()
+            FLAGS.model = str(checkpoint_path.joinpath(latest_ckpt_file))
+        else:
+            FLAGS.model = FLAGS.checkpoint_folder
+        # Default to .pt2; user can specify .pte via -o flag
+        suffix = Path(FLAGS.output).suffix
+        if suffix not in (".pte", ".pt2"):
+            FLAGS.output = str(Path(FLAGS.output).with_suffix(".pt2"))
+        freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
     else:
         raise RuntimeError(
             f"Unsupported command '{FLAGS.command}' for the pt_expt backend."