Merge branch 'master' into 0225-lmdb-dataloader

Author: iProzd

.github/workflows/build_wheel.yml

@@ -187,7 +187,7 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5
 
   pass:
     name: Pass testing build wheels

.github/workflows/test_cc.yml

@@ -98,7 +98,7 @@ jobs:
           ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }}
           ENABLE_PADDLE: ${{ matrix.enable_paddle && '1' || '0' }}
         if: ${{ !matrix.check_memleak }}
-      - uses: codecov/codecov-action@v5
+      - uses: codecov/codecov-action@v6
         with:
           use_oidc: true
     permissions:

.github/workflows/test_python.yml

@@ -76,7 +76,7 @@ jobs:
           name: split-${{ matrix.python }}-${{ matrix.group }}
           path: .test_durations_${{ matrix.group }}
           include-hidden-files: true
-      - uses: codecov/codecov-action@v5
+      - uses: codecov/codecov-action@v6
         with:
           use_oidc: true
     permissions:

.gitignore

@@ -62,6 +62,8 @@ test_dp_test_*.out
 
 # Training and model output files
 *.pth
+*.pte
+*.pt2
 *.ckpt*
 checkpoint
 lcurve.out

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.6
+    rev: v0.15.8
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -62,7 +62,7 @@ repos:
           - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.1
+    rev: v22.1.2
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)
@@ -76,7 +76,7 @@ repos:
         exclude: ^(source/3rdparty|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
-    rev: v3.12.0-2
+    rev: v3.13.0-1
     hooks:
       - id: shfmt
   # CMake

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/atomic_model/base_atomic_model.py

@@ -206,6 +206,20 @@ def change_type_map(
         self.reinit_pair_exclude(
             map_pair_exclude_types(self.pair_exclude_types, remap_index)
         )
+        if has_new_type:
+            xp = array_api_compat.array_namespace(self.out_bias)
+            extend_shape = [
+                self.out_bias.shape[0],
+                len(type_map),
+                *list(self.out_bias.shape[2:]),
+            ]
+            device = array_api_compat.device(self.out_bias)
+            extend_bias = xp.zeros(
+                extend_shape, dtype=self.out_bias.dtype, device=device
+            )
+            self.out_bias = xp.concat([self.out_bias, extend_bias], axis=1)
+            extend_std = xp.ones(extend_shape, dtype=self.out_std.dtype, device=device)
+            self.out_std = xp.concat([self.out_std, extend_std], axis=1)
         self.out_bias = self.out_bias[:, remap_index, :]
         self.out_std = self.out_std[:, remap_index, :]
 

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -17,9 +17,6 @@
     get_multiple_nlist_key,
     nlist_distinguish_types,
 )
-from deepmd.env import (
-    GLOBAL_NP_FLOAT_PRECISION,
-)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -70,6 +67,7 @@ def __init__(
         self,
         models: list[BaseAtomicModel],
         type_map: list[str],
+        weights: str | list[float] = "mean",
         **kwargs: Any,
     ) -> None:
         super().__init__(type_map, **kwargs)
@@ -100,6 +98,15 @@ def __init__(
         self.mapping_list = mapping_list
         assert len(err_msg) == 0, "\n".join(err_msg)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
+        if isinstance(weights, str):
+            assert weights in ["sum", "mean"]
+        elif isinstance(weights, list):
+            assert len(weights) == len(models)
+        else:
+            raise ValueError(
+                f"'weights' must be a string ('sum' or 'mean') or a list of float of length {len(models)}."
+            )
+        self.weights = weights
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -241,7 +248,7 @@ def forward_atomic(
             the result dict, defined by the fitting net output def.
         """
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
-        nframes, nloc, nnei = nlist.shape
+        nframes, _nloc, _nnei = nlist.shape
         extended_coord = xp.reshape(extended_coord, (nframes, -1, 3))
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
@@ -323,20 +330,23 @@ def serialize(self) -> dict:
         dd.update(
             {
                 "@class": "Model",
-                "@version": 2,
+                "@version": 3,
                 "type": "linear",
                 "models": [model.serialize() for model in self.models],
                 "type_map": self.type_map,
+                "weights": self.weights,
             }
         )
         return dd
 
     @classmethod
     def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
         data = data.copy()
-        check_version_compatibility(data.pop("@version", 2), 2, 2)
+        check_version_compatibility(data.pop("@version", 2), 3, 2)
         data.pop("@class", None)
         data.pop("type", None)
+        if "weights" not in data:
+            data["weights"] = "mean"
         models = [
             BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)
             for model in data["models"]
@@ -396,16 +406,33 @@ def _compute_weight(
         nlists_: list[Array],
     ) -> list[Array]:
         """This should be a list of user defined weights that matches the number of models to be combined."""
-        xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlists_)
+        xp = array_api_compat.array_namespace(extended_coord, extended_atype)
         nmodels = len(self.models)
         nframes, nloc, _ = nlists_[0].shape
         dev = array_api_compat.device(extended_coord)
-        # the dtype of weights is the interface data type.
-        return [
-            xp.ones((nframes, nloc, 1), dtype=GLOBAL_NP_FLOAT_PRECISION, device=dev)
-            / nmodels
-            for _ in range(nmodels)
-        ]
+        if isinstance(self.weights, str):
+            if self.weights == "sum":
+                return [
+                    xp.ones((nframes, nloc, 1), dtype=extended_coord.dtype, device=dev)
+                    for _ in range(nmodels)
+                ]
+            elif self.weights == "mean":
+                return [
+                    xp.ones((nframes, nloc, 1), dtype=extended_coord.dtype, device=dev)
+                    / nmodels
+                    for _ in range(nmodels)
+                ]
+            else:
+                raise ValueError(
+                    "`weights` must be 'sum' or 'mean' when provided as a string."
+                )
+        elif isinstance(self.weights, list):
+            return [
+                xp.ones((nframes, nloc, 1), dtype=extended_coord.dtype, device=dev) * w
+                for w in self.weights
+            ]
+        else:
+            raise NotImplementedError
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

deepmd/dpmodel/atomic_model/property_atomic_model.py

@@ -27,7 +27,7 @@ def __init__(
 
     def get_compute_stats_distinguish_types(self) -> bool:
         """Get whether the fitting net computes stats which are not distinguished between different types of atoms."""
-        return False
+        return self.fitting_net.get_distinguish_types()
 
     def get_intensive(self) -> bool:
         """Whether the fitting property is intensive."""
@@ -51,6 +51,10 @@ def apply_out_stat(
 
         """
         out_bias, out_std = self._fetch_out_stat(self.bias_keys)
-        for kk in self.bias_keys:
-            ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]
+        if self.get_compute_stats_distinguish_types():
+            for kk in self.bias_keys:
+                ret[kk] = ret[kk] * out_std[kk][atype] + out_bias[kk][atype]
+        else:
+            for kk in self.bias_keys:
+                ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]
         return ret

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,
@@ -344,6 +345,7 @@ def __init__(
         self.concat_output_tebd = concat_output_tebd
         self.trainable = trainable
         self.precision = precision
+        self.compress = False
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
@@ -535,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,
@@ -557,7 +559,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa1",
-            "@version": 2,
+            "@version": 3 if self.compress else 2,
             "rcut": obj.rcut,
             "rcut_smth": obj.rcut_smth,
             "sel": obj.sel,
@@ -602,20 +604,36 @@ def serialize(self) -> dict:
         }
         if obj.tebd_input_mode in ["strip"]:
             data.update({"embeddings_strip": obj.embeddings_strip.serialize()})
+        if self.compress:
+            compress_dict: dict = {
+                "@variables": {
+                    "type_embd_data": to_numpy_array(self.type_embd_data),
+                },
+                "geo_compress": self.geo_compress,
+            }
+            if self.geo_compress:
+                compress_dict["@variables"]["compress_data"] = [
+                    to_numpy_array(d) for d in self.compress_data
+                ]
+                compress_dict["@variables"]["compress_info"] = [
+                    to_numpy_array(i) for i in self.compress_info
+                ]
+            data["compress"] = compress_dict
         return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptDPA1":
         """Deserialize from dict."""
         data = data.copy()
-        check_version_compatibility(data.pop("@version"), 2, 1)
+        check_version_compatibility(data.pop("@version"), 3, 1)
         data.pop("@class")
         data.pop("type")
         variables = data.pop("@variables")
         embeddings = data.pop("embeddings")
         type_embedding = data.pop("type_embedding")
         attention_layers = data.pop("attention_layers")
         env_mat = data.pop("env_mat")
+        compress = data.pop("compress", None)
         tebd_input_mode = data["tebd_input_mode"]
         if tebd_input_mode in ["strip"]:
             embeddings_strip = data.pop("embeddings_strip")
@@ -637,8 +655,20 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
         obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize(
             attention_layers
         )
+        if compress is not None:
+            obj._load_compress_data(compress)
         return obj
 
+    def _load_compress_data(self, compress: dict) -> None:
+        """Load compression state from serialized data."""
+        variables = compress["@variables"]
+        self.type_embd_data = variables["type_embd_data"]
+        self.geo_compress = compress.get("geo_compress", False)
+        if self.geo_compress:
+            self.compress_data = variables["compress_data"]
+            self.compress_info = variables["compress_info"]
+        self.compress = True
+
     @classmethod
     def update_sel(
         cls,
@@ -1057,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))
@@ -1076,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)
@@ -1084,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)
             )
@@ -1111,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: