feat(dpa4c): add compact invariant descriptor DPA4C 🎉🎉🎉

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -254,6 +254,10 @@ def get_default_chg_spin(self) -> list[float] | None:
         """Get the default charge_spin values."""
         return None
 
+    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
+        """Get the row range each charge_spin value indexes, or None."""
+        return None
+
     def reinit_atom_exclude(
         self,
         exclude_types: list[int] = [],

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -154,6 +154,12 @@ def get_default_chg_spin(self) -> list[float] | None:
             return self.descriptor.get_default_chg_spin()
         return None
 
+    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
+        """Get the row range each charge_spin value indexes, or None."""
+        if self.add_chg_spin_ebd:
+            return self.descriptor.get_chg_spin_table_ranges()
+        return None
+
     def uses_graph_lower(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.uses_graph_lower())
@@ -182,6 +188,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.
 
@@ -731,6 +740,20 @@ def get_default_chg_spin(self) -> "Array | None":
             lambda m: m.get_default_chg_spin(),
         )[1]
 
+    def get_chg_spin_table_ranges(self) -> list[tuple[int, int]] | None:
+        """The shared table row ranges, if the children agree.
+
+        A condition reaches every consuming child, so it must address the
+        tables of all of them. Children that disagree share no acceptable
+        state, which the composition reports as an unconstrained domain
+        rather than silently enforcing one child's tables on the others.
+        """
+        return self._agreed_default(
+            self._chg_spin_consumers(),
+            lambda m: m.get_chg_spin_table_ranges() is not None,
+            lambda m: m.get_chg_spin_table_ranges(),
+        )[1]
+
     def has_default_fparam(self) -> bool:
         """Whether every active child shares one default frame parameter."""
         return self._agreed_default(

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)