Merge branch 'master' into 0304_add_chg_spin

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>

Author: iProzd

deepmd/backend/pt_expt.py

@@ -32,7 +32,7 @@
 class PyTorchExportableBackend(Backend):
     """PyTorch exportable backend."""
 
-    name = "PyTorch Exportable"
+    name = "PyTorch-Exportable"
     """The formal name of the backend."""
     features: ClassVar[Backend.Feature] = (
         Backend.Feature.ENTRY_POINT
@@ -63,7 +63,7 @@ def entry_point_hook(self) -> Callable[["Namespace"], None]:
         Callable[[Namespace], None]
             The entry point hook of the backend.
         """
-        from deepmd.pt.entrypoints.main import main as deepmd_main
+        from deepmd.pt_expt.entrypoints.main import main as deepmd_main
 
         return deepmd_main
 
@@ -76,7 +76,11 @@ def deep_eval(self) -> type["DeepEvalBackend"]:
         type[DeepEvalBackend]
             The Deep Eval backend of the backend.
         """
-        raise NotImplementedError
+        from deepmd.pt_expt.infer.deep_eval import (
+            DeepEval,
+        )
+
+        return DeepEval
 
     @property
     def neighbor_stat(self) -> type["NeighborStat"]:
@@ -87,7 +91,11 @@ def neighbor_stat(self) -> type["NeighborStat"]:
         type[NeighborStat]
             The neighbor statistics of the backend.
         """
-        raise NotImplementedError
+        from deepmd.pt_expt.utils.neighbor_stat import (
+            NeighborStat,
+        )
+
+        return NeighborStat
 
     @property
     def serialize_hook(self) -> Callable[[str], dict]:
@@ -98,7 +106,11 @@ def serialize_hook(self) -> Callable[[str], dict]:
         Callable[[str], dict]
             The serialize hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.pt_expt.utils.serialization import (
+            serialize_from_file,
+        )
+
+        return serialize_from_file
 
     @property
     def deserialize_hook(self) -> Callable[[str, dict], None]:
@@ -109,4 +121,8 @@ def deserialize_hook(self) -> Callable[[str, dict], None]:
         Callable[[str, dict], None]
             The deserialize hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.pt_expt.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        return deserialize_to_file

deepmd/dpmodel/array_api.py

@@ -32,6 +32,15 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     # torch.take_along_dim requires int64 indices
     if array_api_compat.is_torch_array(indices):
         indices = xp.astype(indices, xp.int64)
+    if array_api_compat.is_torch_array(arr):
+        # Use torch.gather directly for torch.export dynamic shape compatibility.
+        # array_api_compat's take_along_axis / torch.take_along_dim specializes
+        # the source dimension size to a constant during torch.export tracing,
+        # breaking dynamic shape export.  torch.gather is the underlying
+        # primitive and handles symbolic shapes correctly.
+        import torch
+
+        return torch.gather(arr, axis, indices)
     if Version(xp.__array_api_version__) >= Version("2024.12"):
         # see: https://github.com/data-apis/array-api-strict/blob/d086c619a58f35c38240592ef994aa19ca7beebc/array_api_strict/_indexing_functions.py#L30-L39
         return xp.take_along_axis(arr, indices, axis=axis)
@@ -62,6 +71,24 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     return xp_swapaxes(out, axis, -1)
 
 
+def xp_take_first_n(arr: Array, dim: int, n: int) -> Array:
+    """Take the first *n* elements along *dim*.
+
+    For torch tensors, uses ``torch.index_select`` so that
+    ``torch.export`` does not emit a contiguity guard that would
+    prevent the ``nall == nloc`` (no-PBC) case from working.
+    For numpy / jax, uses regular slicing.
+    """
+    if array_api_compat.is_torch_array(arr):
+        import torch
+
+        indices = torch.arange(n, dtype=torch.int64, device=arr.device)
+        return torch.index_select(arr, dim, indices)
+    slices = [slice(None)] * arr.ndim
+    slices[dim] = slice(0, n)
+    return arr[tuple(slices)]
+
+
 def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array:
     """Reduces all values from the src tensor to the indices specified in the index tensor.
 

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -13,6 +13,7 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     NativeOP,
@@ -62,6 +63,45 @@ def __init__(
         self.rcond = rcond
         self.preset_out_bias = preset_out_bias
         self.data_stat_protect = data_stat_protect
+        self._observed_type: list[str] | None = None
+
+    @property
+    def observed_type(self) -> list[str] | None:
+        """Get the observed element type list from data statistics."""
+        return self._observed_type
+
+    def _collect_and_set_observed_type(
+        self,
+        sampled_func: Callable[[], list[dict]],
+        stat_file_path: DPPath | None,
+        preset_observed_type: list[str] | None,
+    ) -> None:
+        """Collect observed types with priority: preset > stat_file > compute.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames.
+        stat_file_path
+            The path to the statistics files (should already include type_map suffix).
+        preset_observed_type
+            User-specified observed types that take highest priority.
+        """
+        from deepmd.dpmodel.utils.stat import (
+            _restore_observed_type_from_file,
+            _save_observed_type_to_file,
+            collect_observed_types,
+        )
+
+        if preset_observed_type is not None:
+            self._observed_type = preset_observed_type
+        else:
+            observed = _restore_observed_type_from_file(stat_file_path)
+            if observed is None:
+                sampled = sampled_func()
+                observed = collect_observed_types(sampled, self.type_map)
+                _save_observed_type_to_file(stat_file_path, observed)
+            self._observed_type = observed
 
     def init_out_stat(self) -> None:
         """Initialize the output bias."""
@@ -211,7 +251,7 @@ def forward_common_atomic(
         """
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
         _, nloc, _ = nlist.shape
-        atype = extended_atype[:, :nloc]
+        atype = xp_take_first_n(extended_atype, 1, nloc)
         if self.pair_excl is not None:
             pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype)
             # exclude neighbors in the nlist
@@ -229,7 +269,7 @@ def forward_common_atomic(
         ret_dict = self.apply_out_stat(ret_dict, atype)
 
         # nf x nloc
-        atom_mask = ext_atom_mask[:, :nloc]
+        atom_mask = xp_take_first_n(ext_atom_mask, 1, nloc)
         if self.atom_excl is not None:
             atom_mask = xp.logical_and(
                 atom_mask, self.atom_excl.build_type_exclude_mask(atype)
@@ -271,6 +311,29 @@ 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 True
 
+    def compute_or_load_stat(
+        self,
+        sampled_func: Callable[[], list[dict]],
+        stat_file_path: DPPath | None = None,
+        compute_or_load_out_stat: bool = True,
+        preset_observed_type: list[str] | None = None,
+    ) -> None:
+        """Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+
+        Parameters
+        ----------
+        sampled_func
+            The lazy sampled function to get data frames from different data systems.
+        stat_file_path
+            The path to the stat file.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics
+            (e.g. mean and standard deviation of descriptors).
+        """
+        raise NotImplementedError
+
     def compute_or_load_out_stat(
         self,
         merged: Callable[[], list[dict]] | list[dict],
@@ -332,19 +395,19 @@ def wrapped_sampler() -> list[dict]:
                 atom_exclude_types = self.atom_excl.get_exclude_types()
                 for sample in sampled:
                     sample["atom_exclude_types"] = list(atom_exclude_types)
-            if (
-                "find_fparam" not in sampled[0]
-                and "fparam" not in sampled[0]
-                and self.has_default_fparam()
-            ):
+            # For systems where fparam is missing (find_fparam == 0),
+            # fill with default fparam if available and mark as found.
+            if self.has_default_fparam():
                 default_fparam = self.get_default_fparam()
                 if default_fparam is not None:
                     default_fparam_np = np.array(default_fparam)
                     for sample in sampled:
-                        nframe = sample["atype"].shape[0]
-                        sample["fparam"] = np.tile(
-                            default_fparam_np.reshape(1, -1), (nframe, 1)
-                        )
+                        if "find_fparam" in sample and not sample["find_fparam"]:
+                            nframe = sample["atype"].shape[0]
+                            sample["fparam"] = np.tile(
+                                default_fparam_np.reshape(1, -1), (nframe, 1)
+                            )
+                            sample["find_fparam"] = np.bool_(True)
             return sampled
 
         return wrapped_sampler

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -8,6 +8,7 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
@@ -181,7 +182,7 @@ def forward_atomic(
 
         """
         nframes, nloc, nnei = nlist.shape
-        atype = extended_atype[:, :nloc]
+        atype = xp_take_first_n(extended_atype, 1, nloc)
 
         if self.fitting_net.get_dim_fparam() > 0 and fparam is None:
             # use default fparam
@@ -221,6 +222,7 @@ def compute_or_load_stat(
         sampled_func: Callable[[], list[dict]],
         stat_file_path: DPPath | None = None,
         compute_or_load_out_stat: bool = True,
+        preset_observed_type: list[str] | None = None,
     ) -> None:
         """Compute or load the statistics parameters of the model,
         such as mean and standard deviation of descriptors or the energy bias of the fitting net.
@@ -247,6 +249,10 @@ def compute_or_load_stat(
         if compute_or_load_out_stat:
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
+        self._collect_and_set_observed_type(
+            wrapped_sampler, stat_file_path, preset_observed_type
+        )
+
     def change_type_map(
         self, type_map: list[str], model_with_new_type_stat: Any | None = None
     ) -> None:

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -349,6 +349,7 @@ def compute_or_load_stat(
         sampled_func: Callable[[], list[dict]],
         stat_file_path: DPPath | None = None,
         compute_or_load_out_stat: bool = True,
+        preset_observed_type: list[str] | None = None,
     ) -> None:
         """Compute or load the statistics parameters of the model.
 
@@ -364,9 +365,21 @@ def compute_or_load_stat(
         compute_or_load_out_stat : bool
             Whether to compute the output statistics.
         """
+        # Compute observed type once at parent level, then propagate to
+        # sub-models via preset_observed_type to avoid redundant computation.
+        obs_stat_path = stat_file_path
+        if obs_stat_path is not None and self.type_map is not None:
+            obs_stat_path = obs_stat_path / " ".join(self.type_map)
+        self._collect_and_set_observed_type(
+            sampled_func, obs_stat_path, preset_observed_type
+        )
+
         for md in self.models:
             md.compute_or_load_stat(
-                sampled_func, stat_file_path, compute_or_load_out_stat=False
+                sampled_func,
+                stat_file_path,
+                compute_or_load_out_stat=False,
+                preset_observed_type=self._observed_type,
             )
 
         if stat_file_path is not None and self.type_map is not None:

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -216,6 +216,7 @@ def compute_or_load_stat(
         sampled_func: Callable[[], list[dict]],
         stat_file_path: DPPath | None = None,
         compute_or_load_out_stat: bool = True,
+        preset_observed_type: list[str] | None = None,
     ) -> None:
         """Compute or load the statistics parameters of the model.
 
@@ -235,6 +236,15 @@ def compute_or_load_stat(
             wrapped_sampler = self._make_wrapped_sampler(sampled_func)
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
+        if stat_file_path is not None and self.type_map is not None:
+            stat_file_path /= " ".join(self.type_map)
+
+        self._collect_and_set_observed_type(
+            sampled_func if callable(sampled_func) else lambda: sampled_func,
+            stat_file_path,
+            preset_observed_type,
+        )
+
     def forward_atomic(
         self,
         extended_coord: Array,

deepmd/dpmodel/descriptor/dpa1.py

@@ -242,6 +242,8 @@ class DescrptDPA1(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2208.08236.
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         rcut: float,
@@ -663,7 +665,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        min_nbor_dist, sel = UpdateSel().update_one_sel(
+        min_nbor_dist, sel = cls._update_sel_cls().update_one_sel(
             train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
         )
         local_jdata_cpy["sel"] = sel[0]

deepmd/dpmodel/descriptor/dpa2.py

@@ -441,6 +441,8 @@ class DescrptDPA2(NativeOP, BaseDescriptor):
        Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -1115,7 +1117,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repinit_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/dpa3.py

@@ -338,6 +338,8 @@ class DescrptDPA3(NativeOP, BaseDescriptor):
        arXiv preprint arXiv:2506.01686 (2025).
     """
 
+    _update_sel_cls = UpdateSel
+
     def __init__(
         self,
         ntypes: int,
@@ -799,7 +801,7 @@ def update_sel(
             The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        update_sel = UpdateSel()
+        update_sel = cls._update_sel_cls()
         min_nbor_dist, repflow_e_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -51,7 +51,7 @@ class BD(ABC, PluginVariant, make_plugin_registry("descriptor")):
         def __new__(cls, *args: Any, **kwargs: Any) -> Any:
             if cls is BD:
                 cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__))
-            return super().__new__(cls)
+            return object.__new__(cls)
 
         @abstractmethod
         def get_rcut(self) -> float: