Merge remote-tracking branch 'upstream/master' into pr/install-skill

Author: OutisLi

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -119,7 +119,7 @@ def __init__(
         self.fitting_net = fitting
         self.fitting_net.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
-        self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0
+        self.add_chg_spin_ebd: bool = self.descriptor.has_chg_spin_ebd()
         # Structural capability: only descriptors with a native spin
         # conditioning mechanism (currently DPA4) accept a ``spin`` kwarg on
         # ``call_graph`` at all -- unlike ``charge_spin``, which every
@@ -182,6 +182,10 @@ def supports_graph_export(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.supports_graph_export())
 
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Delegates to this model's own descriptor."""
+        return bool(self.descriptor.compression_needs_min_nbor_dist())
+
     def supports_native_spin(self) -> bool:
         """Delegates to this model's own descriptor (cached at construction)."""
         return self._supports_native_spin

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -303,6 +303,15 @@ def enable_compression(
                 check_frequency,
             )
 
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Required as soon as ANY child consumes it.
+
+        The statistic is measured once and handed to every child, so a single
+        child that tabulates from the shortest observed distance keeps the
+        neighbor-statistics pass for the whole composition.
+        """
+        return any(m.compression_needs_min_nbor_dist() for m in self.models)
+
     def uses_graph_lower(self) -> bool:
         """Graph-capable iff EVERY child supports the graph lower.
 

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -190,6 +190,17 @@ def enable_compression(
             """
             raise NotImplementedError("This atomi model doesn't support compression!")
 
+        def compression_needs_min_nbor_dist(self) -> bool:
+            """Whether :meth:`enable_compression` consumes ``min_nbor_dist``.
+
+            Returns
+            -------
+            bool
+                Concrete default ``True``, so a model that does not report
+                otherwise keeps the neighbor-statistics pass.
+            """
+            return True
+
         def make_atom_mask(
             self,
             atype: t_tensor,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -505,3 +505,14 @@ def enable_compression(
     ) -> None:
         """Pairtab model does not support compression."""
         pass
+
+    def compression_needs_min_nbor_dist(self) -> bool:
+        """Return whether compression consumes the minimum neighbor distance.
+
+        Returns
+        -------
+        bool
+            Always ``False``. The tabulated pair potential carries its own
+            domain, so compression is a no-op here.
+        """
+        return False

deepmd/dpmodel/descriptor/__init__.py

@@ -11,6 +11,9 @@
 from .dpa4 import (
     DescrptDPA4,
 )
+from .dpa4c import (
+    DescrptDPA4C,
+)
 from .hybrid import (
     DescrptHybrid,
 )
@@ -38,6 +41,7 @@
     "DescrptDPA2",
     "DescrptDPA3",
     "DescrptDPA4",
+    "DescrptDPA4C",
     "DescrptHybrid",
     "DescrptSeA",
     "DescrptSeAttenV2",

deepmd/dpmodel/descriptor/dpa3.py

@@ -33,6 +33,12 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.charge_state import (
+    CHARGE_OFFSET,
+    CHARGE_TABLE_ROWS,
+    MULTIPLICITY_TABLE_ROWS,
+    validate_charge_state,
+)
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
 )
@@ -468,11 +474,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
 
         self.use_econf_tebd = use_econf_tebd
         self.add_chg_spin_ebd = add_chg_spin_ebd
-        if default_chg_spin is not None and len(default_chg_spin) != 2:
-            raise ValueError(
-                "default_chg_spin must have exactly 2 values [charge, spin]"
-            )
-        self.default_chg_spin = default_chg_spin
+        self.default_chg_spin = (
+            None
+            if default_chg_spin is None
+            else validate_charge_state(default_chg_spin)
+        )
         self.use_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
@@ -494,18 +500,16 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
 
         if self.add_chg_spin_ebd:
             self.cs_activation_fn = get_activation_fn(activation_function)
-            # -100 ~ 100 is a conservative bound
             self.chg_embedding = TypeEmbedNet(
-                ntypes=200,
+                ntypes=CHARGE_TABLE_ROWS,
                 neuron=[self.tebd_dim],
                 padding=True,
                 activation_function="Linear",
                 precision=precision,
                 seed=child_seed(seed, 3),
             )
-            # 100 is a conservative upper bound
             self.spin_embedding = TypeEmbedNet(
-                ntypes=100,
+                ntypes=MULTIPLICITY_TABLE_ROWS,
                 neuron=[self.tebd_dim],
                 padding=True,
                 activation_function="Linear",
@@ -543,6 +547,10 @@ def get_dim_chg_spin(self) -> int:
         """Returns the dimension of charge_spin input."""
         return 2 if self.add_chg_spin_ebd else 0
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Return whether a frame charge/spin condition is configured."""
+        return self.add_chg_spin_ebd
+
     def get_default_chg_spin(self) -> list[float] | None:
         """Returns the default charge_spin values."""
         return self.default_chg_spin
@@ -755,7 +763,7 @@ def call(
             assert self.spin_embedding is not None
             chg_tebd = self.chg_embedding.call()
             spin_tebd = self.spin_embedding.call()
-            charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+            charge = xp.astype(charge_spin[:, 0], xp.int64) + CHARGE_OFFSET
             spin = xp.astype(charge_spin[:, 1], xp.int64)
             chg_ebd = xp.reshape(
                 xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0),

deepmd/dpmodel/descriptor/dpa4.py

@@ -73,6 +73,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.charge_state import (
+    validate_charge_state,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -784,10 +787,10 @@ def __init__(
         self.edge_cartesian = bool(edge_cartesian)
         self.node_cartesian = str(node_cartesian)
         self.add_chg_spin_ebd = bool(add_chg_spin_ebd)
-        if default_chg_spin is not None and len(default_chg_spin) != 2:
-            raise ValueError("`default_chg_spin` must contain [charge, spin].")
         self.default_chg_spin = (
-            None if default_chg_spin is None else [float(x) for x in default_chg_spin]
+            None
+            if default_chg_spin is None
+            else validate_charge_state(default_chg_spin)
         )
 
         # === Native per-atom spin embedding ===
@@ -2282,6 +2285,10 @@ def get_ntypes(self) -> int:
     def get_type_map(self) -> list[str]:
         return self.type_map if self.type_map is not None else []
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Return whether a frame charge/spin condition is configured."""
+        return self.charge_spin_embedding is not None
+
     def get_dim_chg_spin(self) -> int:
         """Return the charge/spin condition width."""
         return 2 if self.add_chg_spin_ebd else 0

deepmd/dpmodel/descriptor/dpa4_nn/__init__.py

@@ -73,6 +73,10 @@
     merge_lora_into_base,
     strip_lora_from_extra_state,
 )
+from .mlp import (
+    SwiGLUMLP,
+    resolve_swiglu_hidden_width,
+)
 from .norm import (
     EquivariantRMSNorm,
     ReducedEquivariantRMSNorm,
@@ -159,6 +163,7 @@
     "SeZMTypeEmbedding",
     "SpinEmbedding",
     "SwiGLU",
+    "SwiGLUMLP",
     "WignerDCalculator",
     "apply_lora_to_sezm",
     "build_cartesian_basis",
@@ -189,6 +194,7 @@
     "quaternion_z_rotation",
     "resolve_s2_grid_resolution",
     "resolve_so3_grid",
+    "resolve_swiglu_hidden_width",
     "safe_norm",
     "segment_envelope_gated_softmax",
     "so3_packed_index",

deepmd/dpmodel/descriptor/dpa4_nn/embedding.py

@@ -45,6 +45,11 @@
 from deepmd.dpmodel.utils.type_embed import (
     remap_atype_to_padding,
 )
+from deepmd.utils.charge_state import (
+    CHARGE_OFFSET,
+    CHARGE_TABLE_ROWS,
+    MULTIPLICITY_TABLE_ROWS,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -128,26 +133,41 @@ def __init__(
         # === Step 2. Register the embedding table parameter ===
         self.adam_type_embedding = table
 
-    def call(self, atype: Any) -> Any:
+    def call(self, atype: Any | None = None) -> Any:
         """
         Gather type embeddings.
 
         Parameters
         ----------
         atype
-            Atom types with shape (...,). Valid type range is [0, ntypes-1].
+            Atom types with shape (...). Valid type range is [0, ntypes-1].
+            If omitted, return the complete embedding table, including the
+            optional padding row. This form is used by graph-native descriptor
+            ABIs that precompute the table once per forward call.
 
         Returns
         -------
         Array
-            Type embeddings with shape (..., embed_dim).
+            Gathered type embeddings with shape ``(..., embed_dim)`` when
+            ``atype`` is provided. Otherwise, the complete table with shape
+            ``(ntypes + int(padding), embed_dim)``.
         """
+        # === Step 1. Return the complete graph-native lookup table ===
+        if atype is None:
+            xp = array_api_compat.array_namespace(self.adam_type_embedding)
+            return xp_asarray_nodetach(
+                xp,
+                self.adam_type_embedding[...],
+                device=array_api_compat.device(self.adam_type_embedding),
+            )
+
+        # === Step 2. Gather rows for an explicit atom-type tensor ===
         xp = array_api_compat.array_namespace(atype)
         weight = xp_asarray_nodetach(
             xp, self.adam_type_embedding[...], device=array_api_compat.device(atype)
         )
-        # torch.embedding gather: flatten the indices to int64, take the rows,
-        # then restore the original index shape.
+        # Flattening provides one backend-neutral gather while preserving every
+        # leading batch or graph dimension on restoration.
         index = xp.astype(xp.reshape(atype, (-1,)), xp.int64)
         if self.padding:
             index = remap_atype_to_padding(index, self.ntypes + 1)
@@ -869,15 +889,15 @@ def __init__(
             raise ValueError("`embed_dim` must be positive")
 
         self.charge_embedding = SeZMTypeEmbedding(
-            ntypes=200,
+            ntypes=CHARGE_TABLE_ROWS,
             embed_dim=self.embed_dim,
             precision=self.precision,
             seed=child_seed(seed, 0),
             trainable=self.trainable,
             padding=False,
         )
         self.spin_embedding = SeZMTypeEmbedding(
-            ntypes=100,
+            ntypes=MULTIPLICITY_TABLE_ROWS,
             embed_dim=self.embed_dim,
             precision=self.precision,
             seed=child_seed(seed, 1),
@@ -908,7 +928,7 @@ def call(self, charge_spin: Any) -> Any:
             Mixed condition embedding with shape (nf, embed_dim).
         """
         xp = array_api_compat.array_namespace(charge_spin)
-        charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+        charge = xp.astype(charge_spin[:, 0], xp.int64) + CHARGE_OFFSET
         spin = xp.astype(charge_spin[:, 1], xp.int64)
         charge_embed = self.charge_embedding(charge)
         spin_embed = self.spin_embedding(spin)

deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py

@@ -99,6 +99,37 @@ def _build_frame_degree_index(
     raise ValueError("`coefficient_layout` must be either 'packed' or 'm_major'")
 
 
+def _degree_batched_matmul(xp: Any, coeff: Any, weight: Any) -> Any:
+    """Contract ``einsum("ndfi,dio->ndfo")`` batched over the degree axis.
+
+    Parameters
+    ----------
+    xp : Any
+        The array namespace of ``coeff``.
+    coeff : Array
+        Coefficients with shape ``(N, D, F, i)``.
+    weight : Array
+        Per-degree weights with shape ``(D, i, o)``.
+
+    Returns
+    -------
+    Array
+        Contracted coefficients with shape ``(N, D, F, o)``.
+
+    Notes
+    -----
+    Batching over the ``(D, F)`` axes, not over ``N``: expanding ``weight``
+    across ``F`` costs ``D*F*i*o`` elements, whereas batching over ``N``
+    (or collapsing ``N*F``, which needs a materialized permuted copy of
+    ``coeff``) touches ``N*D*F*i`` elements — a factor ``N/o`` more. No
+    reshape is involved, so an empty ``N`` batch (empty graph/edge set, or
+    a distributed rank owning no nodes) flows through naturally.
+    """
+    coeff_df = xp.permute_dims(coeff, (1, 2, 0, 3))  # (D, F, N, i)
+    out = xp.matmul(coeff_df, weight[:, None, :, :])  # (D, F, N, o)
+    return xp.permute_dims(out, (2, 0, 1, 3))  # (N, D, F, o)
+
+
 def _project_frames(coeff: Any, proj: ChannelLinear, n_frames: int) -> Any:
     """
     Apply a channel-only linear map to each Wigner-D frame independently.
@@ -493,9 +524,8 @@ def call(self, coeff: Any) -> Any:
         weight = xp_asarray_nodetach(xp, self.weight[...], device=device)
         degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device)
         weight = xp.take(weight, degree_index, axis=0)
-        # einsum "ndfi,dio->ndfo" as a broadcast batched matmul:
-        # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o)
-        return xp.matmul(coeff, weight[None, ...])
+        # Batched over the (D, F) axes, never over N -- see the helper's note.
+        return _degree_batched_matmul(xp, coeff, weight)
 
     def serialize(self) -> dict[str, Any]:
         """Serialize the FrameContract to a dict."""
@@ -575,9 +605,8 @@ def call(self, coeff: Any) -> Any:
         weight = xp_asarray_nodetach(xp, self.weight[...], device=device)
         degree_index = xp_asarray_nodetach(xp, self.degree_index, device=device)
         weight = xp.take(weight, degree_index, axis=0)
-        # einsum "ndfi,dio->ndfo" as a broadcast batched matmul:
-        # (N, D, F, i) @ (1, D, i, o) -> (N, D, F, o)
-        return xp.matmul(coeff, weight[None, ...])
+        # Batched over the (D, F) axes, never over N -- see the helper's note.
+        return _degree_batched_matmul(xp, coeff, weight)
 
     def serialize(self) -> dict[str, Any]:
         """Serialize the FrameExpand to a dict."""