feat(pt-expt): add compact descriptor DPA4C 🎉🎉🎉

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -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/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

@@ -128,26 +128,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)

deepmd/dpmodel/descriptor/dpa4_nn/mlp.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Bias-free SwiGLU multilayer perceptrons for DPA4-family descriptors."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+import numpy as np
+
+from deepmd.dpmodel import (
+    DEFAULT_PRECISION,
+    PRECISION_DICT,
+    NativeOP,
+)
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+from deepmd.dpmodel.utils.network import (
+    NativeLayer,
+)
+from deepmd.dpmodel.utils.seed import (
+    child_seed,
+)
+from deepmd.utils.version import (
+    check_version_compatibility,
+)
+
+from .activation import (
+    SwiGLU,
+)
+
+
+def resolve_swiglu_hidden_width(width: int, multiple: int = 8) -> int:
+    r"""Return the parameter-matched SwiGLU hidden width.
+
+    The post-gate width is :math:`8d/3`, rounded up to ``multiple``. The
+    corresponding hidden affine map produces twice this width for the value
+    and gate branches.
+
+    Parameters
+    ----------
+    width
+        Input and output model width.
+    multiple
+        Alignment multiple for the post-gate hidden width.
+
+    Returns
+    -------
+    int
+        Aligned post-gate hidden width.
+    """
+    if width <= 0:
+        raise ValueError(f"`width` must be positive, got {width}")
+    if multiple <= 0:
+        raise ValueError(f"`multiple` must be positive, got {multiple}")
+    numerator = 8 * int(width)
+    denominator = 3 * int(multiple)
+    return int(multiple) * ((numerator + denominator - 1) // denominator)
+
+
+class SwiGLUMLP(NativeOP):
+    """Apply bias-free SwiGLU hidden layers and a linear output projection.
+
+    For hidden width ``H``, each hidden affine map produces ``2H`` channels.
+    :class:`SwiGLU` splits them into equal gate and value branches and returns
+    ``SiLU(gate) * value`` with width ``H``. The final layer is linear.
+
+    Parameters
+    ----------
+    mlp_layers
+        Layer widths including input, hidden, and output dimensions.
+    output_scale
+        Fixed multiplier applied to the final output.
+    precision
+        Parameter precision.
+    trainable
+        Whether the linear weights are trainable.
+    seed
+        Random seed.
+    """
+
+    def __init__(
+        self,
+        mlp_layers: list[int],
+        *,
+        output_scale: float = 1.0,
+        precision: str = DEFAULT_PRECISION,
+        trainable: bool = True,
+        seed: int | list[int] | None = None,
+    ) -> None:
+        if len(mlp_layers) < 2:
+            raise ValueError("`mlp_layers` must contain input and output widths")
+        if any(width <= 0 for width in mlp_layers):
+            raise ValueError(f"`mlp_layers` must be positive, got {mlp_layers}")
+        self.mlp_layers = [int(width) for width in mlp_layers]
+        self.output_scale = float(output_scale)
+        self.precision = str(precision)
+        self.trainable = bool(trainable)
+
+        layers = []
+        for index, (width_in, width_out) in enumerate(
+            zip(self.mlp_layers[:-1], self.mlp_layers[1:], strict=True)
+        ):
+            is_output = index == len(self.mlp_layers) - 2
+            layers.append(
+                NativeLayer(
+                    width_in,
+                    width_out if is_output else 2 * width_out,
+                    bias=False,
+                    precision=self.precision,
+                    seed=child_seed(seed, index),
+                    trainable=self.trainable,
+                )
+            )
+        self.layers = layers
+        self.activation = SwiGLU()
+
+    def call(self, inputs: Any) -> Any:
+        """Evaluate the SwiGLU MLP.
+
+        Parameters
+        ----------
+        inputs
+            Input with shape ``(..., mlp_layers[0])``.
+
+        Returns
+        -------
+        Any
+            Output with shape ``(..., mlp_layers[-1])``.
+        """
+        return self.call_output(self.call_hidden(inputs))
+
+    def call_hidden(self, inputs: Any) -> Any:
+        """Evaluate every hidden layer and return the latent state.
+
+        The latent state is exposed separately so that several output heads
+        can branch off one trunk evaluation.
+
+        Parameters
+        ----------
+        inputs
+            Input with shape ``(..., mlp_layers[0])``.
+
+        Returns
+        -------
+        Any
+            Activated latent state with shape ``(..., mlp_layers[-2])``.
+        """
+        output = inputs
+        for layer in self.layers[:-1]:
+            output = self.activation(layer(output))
+        return output
+
+    def call_hidden_affine(self, inputs: Any) -> Any:
+        """Apply the first hidden affine map without its activation.
+
+        The affine map is linear and bias free, so an additive shift of the
+        input appears here as an additive shift of the pre-activation. A
+        caller that evaluates the same trunk under several such shifts can
+        therefore share this projection and add each shift afterwards,
+        instead of duplicating the input over the shift axis.
+
+        Parameters
+        ----------
+        inputs
+            Input with shape ``(..., mlp_layers[0])``.
+
+        Returns
+        -------
+        Any
+            Pre-activation with shape ``(..., 2 * mlp_layers[1])``.
+        """
+        return self.layers[0](inputs)
+
+    def call_from_hidden_affine(self, pre_activation: Any) -> Any:
+        """Complete the MLP from the first hidden affine pre-activation.
+
+        Composing this with :meth:`call_hidden_affine` reproduces
+        :meth:`call` exactly, for any number of hidden layers.
+
+        Parameters
+        ----------
+        pre_activation
+            Pre-activation with shape ``(..., 2 * mlp_layers[1])``, as
+            returned by :meth:`call_hidden_affine`.
+
+        Returns
+        -------
+        Any
+            Output with shape ``(..., mlp_layers[-1])``.
+        """
+        output = self.activation(pre_activation)
+        for layer in self.layers[1:-1]:
+            output = self.activation(layer(output))
+        return self.call_output(output)
+
+    def call_output(self, hidden: Any) -> Any:
+        """Apply the final scaled linear projection to a latent state.
+
+        Parameters
+        ----------
+        hidden
+            Latent state with shape ``(..., mlp_layers[-2])``, as returned by
+            :meth:`call_hidden`.
+
+        Returns
+        -------
+        Any
+            Output with shape ``(..., mlp_layers[-1])``.
+        """
+        return self.layers[-1](hidden) * self.output_scale
+
+    def serialize(self) -> dict[str, Any]:
+        """Serialize the MLP configuration and linear weights."""
+        return {
+            "@class": "SwiGLUMLP",
+            "@version": 1,
+            "mlp_layers": self.mlp_layers.copy(),
+            "output_scale": self.output_scale,
+            "precision": np.dtype(PRECISION_DICT[self.precision]).name,
+            "trainable": self.trainable,
+            "@variables": {
+                f"{index}.matrix": to_numpy_array(layer.w)
+                for index, layer in enumerate(self.layers)
+            },
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any]) -> SwiGLUMLP:
+        """Deserialize a :class:`SwiGLUMLP`."""
+        data = data.copy()
+        check_version_compatibility(data.pop("@version"), 1, 1)
+        if data.pop("@class") != "SwiGLUMLP":
+            raise ValueError("Invalid serialized class for SwiGLUMLP")
+        variables = data.pop("@variables")
+        obj = cls(**data)
+        dtype = PRECISION_DICT[obj.precision]
+        for key, value in variables.items():
+            index, _, name = key.partition(".")
+            if name != "matrix":
+                raise ValueError(f"Invalid SwiGLUMLP variable {key!r}")
+            obj.layers[int(index)].w = np.asarray(value, dtype=dtype)
+        return obj