feat(dpa4): DPA4/SeZM on the NeighborGraph lower — graph-native port, multi-rank, native spin, charge-spin/bridging

deepmd/dpmodel/array_api.py

@@ -334,6 +334,48 @@ def xp_setitem_at(x: Array, mask: Array, values: Array) -> Array:
     return x
 
 
+def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> Array:
+    """Draw ``size`` uniform samples in ``[low, high)`` on ``like``'s device.
+
+    Each backend uses its own generator: torch draws with ``torch.rand`` (as
+    pt does, so ``setup_seed`` replays it, with no host copy -- and a host
+    draw would freeze to a constant under tracing); other backends use
+    :mod:`deepmd.utils.random`, which ``setup_seed`` also seeds. Draws are
+    therefore not comparable across backends -- use only for a per-forward
+    random stream, never where a parity test looks.
+
+    Parameters
+    ----------
+    like : Array
+        Reference array supplying backend, dtype and device.
+    size : int
+        Number of samples to draw.
+    low : float
+        Lower bound of the interval.
+    high : float
+        Upper bound of the interval, exclusive.
+
+    Returns
+    -------
+    Array
+        Samples of shape ``(size,)`` matching ``like``.
+    """
+    if array_api_compat.is_torch_array(like):
+        import torch
+
+        return (
+            torch.rand(size, dtype=like.dtype, device=like.device) * (high - low) + low
+        )
+    from deepmd.utils import random as dp_random
+
+    xp = array_api_compat.array_namespace(like)
+    drawn = np.asarray(dp_random.random(size)) * (high - low) + low
+    return xp.astype(
+        xp_asarray_nodetach(xp, drawn, device=array_api_compat.device(like)),
+        like.dtype,
+    )
+
+
 def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> Array:
     """Counts the number of occurrences of each value in x."""
     xp = array_api_compat.array_namespace(x)

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -167,6 +167,29 @@ def has_default_fparam(self) -> bool:
         """Check if the model has default frame parameters."""
         return False
 
+    def uses_graph_lower(self) -> bool:
+        """Returns whether this atomic model supports the NeighborGraph lower.
+
+        Generic capability (concrete default ``False``): the model layer
+        consults it for graph-route eligibility without assuming anything
+        about the atomic model's internal architecture. Implementations
+        answer from their own structure (e.g. a descriptor+fitting model
+        delegates to its descriptor; a composition supports it iff ALL its
+        children do).
+        """
+        return False
+
+    def supports_native_spin(self) -> bool:
+        """Returns whether this atomic model consumes a per-atom spin input.
+
+        Generic capability (concrete default ``False``), the twin of
+        :meth:`uses_graph_lower`: the model layer asks the atomic model
+        directly instead of reaching into it for a descriptor, so the answer
+        stays correct for architectures with no descriptor at all (analytical
+        terms) and for compositions.
+        """
+        return False
+
     def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return None
@@ -381,6 +404,7 @@ def forward_common_atomic_graph(
         fparam: Array | None = None,
         aparam: Array | None = None,
         charge_spin: Array | None = None,
+        spin: Array | None = None,
         comm_dict: dict | None = None,
     ) -> dict:
         """Graph analogue of :meth:`forward_common_atomic` on the flat node axis.
@@ -405,8 +429,14 @@ def forward_common_atomic_graph(
         aparam
             atomic parameter. N x nda
         charge_spin
-            charge/spin conditioning. Unused by the dpa1 graph path; accepted so
-            the interface stays stable for charge/spin-conditioned descriptors.
+            frame-level charge/spin conditioning, forwarded unchanged to
+            :meth:`forward_atomic_graph`, which only passes it on to the
+            descriptor's ``call_graph`` for descriptors that declare
+            ``supports_charge_spin`` (currently DPA4 only).
+        spin
+            flat (N, 3) per-node spin, forwarded unchanged to
+            :meth:`forward_atomic_graph` (and, from there, the descriptor's
+            ``call_graph``); None for spin-less models.
         comm_dict
             MPI communication metadata forwarded to :meth:`forward_atomic_graph`
             (and, from there, the descriptor's ``call_graph``). ``None`` for
@@ -426,6 +456,7 @@ def forward_common_atomic_graph(
             fparam=fparam,
             aparam=aparam,
             charge_spin=charge_spin,
+            spin=spin,
             comm_dict=comm_dict,
         )
         return self._finalize_atomic_ret(ret_dict, output_mask, atype)

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -123,6 +123,22 @@ def __init__(
         self.add_chg_spin_ebd: bool = getattr(
             self.descriptor, "add_chg_spin_ebd", False
         )
+        # 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
+        # descriptor's dense ``call()`` accepts (and ignores) for interface
+        # stability, the graph-native ``call_graph`` signature is
+        # per-descriptor, so ``forward_atomic_graph`` must not pass the
+        # keyword to a descriptor whose ``call_graph`` does not declare it
+        # (that would be a ``TypeError``, not a no-op). Queried via the
+        # ``supports_native_spin`` capability method declared on
+        # ``BaseDescriptor`` (concrete default ``False``; DPA4 overrides).
+        self._supports_native_spin: bool = self.descriptor.supports_native_spin()
+        # Same capability method as ``supports_native_spin`` above, for the
+        # frame-level ``charge_spin`` FiLM kwarg: only DPA4's ``call_graph``
+        # declares it; other descriptors' ``call_graph`` would ``TypeError``
+        # on an unconditional ``charge_spin=`` kwarg.
+        self.supports_charge_spin: bool = self.descriptor.supports_charge_spin()
         super().init_out_stat()
 
     def has_chg_spin_ebd(self) -> bool:
@@ -147,6 +163,14 @@ def get_default_chg_spin(self) -> list[float] | None:
             return self.descriptor.get_default_chg_spin()
         return None
 
+    def uses_graph_lower(self) -> bool:
+        """Delegates to this model's own descriptor."""
+        return bool(self.descriptor.uses_graph_lower())
+
+    def supports_native_spin(self) -> bool:
+        """Delegates to this model's own descriptor (cached at construction)."""
+        return self._supports_native_spin
+
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return self.fitting_net.output_def()
@@ -301,6 +325,7 @@ def forward_atomic_graph(
         fparam: Array | None = None,
         aparam: Array | None = None,
         charge_spin: Array | None = None,
+        spin: Array | None = None,
         comm_dict: dict | None = None,
     ) -> dict[str, Array]:
         """Graph analogue of :meth:`forward_atomic` on the flat node axis.
@@ -321,8 +346,14 @@ def forward_atomic_graph(
         aparam
             atomic parameter. N x nda
         charge_spin
-            charge/spin conditioning. Unused by the dpa1 graph path; accepted so
-            the interface stays stable for charge/spin-conditioned descriptors.
+            frame-level charge/spin conditioning, forwarded to the
+            descriptor's ``call_graph`` only when
+            ``self.supports_charge_spin`` (currently DPA4 only); ignored (not
+            forwarded, never a ``TypeError``) for descriptors without that
+            capability, keeping the interface stable for all of them.
+        spin
+            flat (N, 3) per-node spin, forwarded to the descriptor's
+            ``call_graph``; None for spin-less models.
         comm_dict
             MPI communication metadata forwarded to the descriptor's
             ``call_graph`` (the message-passing part). ``None`` for
@@ -342,9 +373,24 @@ def forward_atomic_graph(
         )
 
         xp = array_api_compat.array_namespace(graph.edge_vec)
-        type_embedding = self.descriptor.type_embedding.call()
+        # Descriptor-owned: dpa1/dpa2 hand out their full tebd table; DPA4
+        # embeds types internally from ``atype`` and returns None.
+        type_embedding = self.descriptor.graph_type_embedding_table()
+        # Only forward the ``spin``/``charge_spin`` keyword to descriptors
+        # whose ``call_graph`` declares it.  Queried through the public
+        # capability -- an override must reach THIS call site too, so the
+        # cached field stays behind the method.
+        spin_kwargs = {"spin": spin} if self.supports_native_spin() else {}
+        charge_spin_kwargs = (
+            {"charge_spin": charge_spin} if self.supports_charge_spin else {}
+        )
         gg, rot_mat = self.descriptor.call_graph(
-            graph, atype, type_embedding=type_embedding, comm_dict=comm_dict
+            graph,
+            atype,
+            type_embedding=type_embedding,
+            comm_dict=comm_dict,
+            **spin_kwargs,
+            **charge_spin_kwargs,
         )
         fparam_node = None
         if fparam is not None: