feat(dpa4): wire charge-spin FiLM + SFPG bridging onto the graph route

Both features already worked inside the shared _call_graph_impl trunk;
only three wiring gaps kept them off the NeighborGraph lower:

- DescrptDPA4.call_graph dropped charge_spin before it reached the
  trunk. It now accepts charge_spin, canonicalizes it (same
  _canonicalize_charge_spin the dense call() adapter uses, so
  call_graph is the single owner of that step on the graph route),
  recovers nf from graph.n_node.shape[0] (static shape, export-safe),
  and threads it through _call_graph_common. Gated end-to-end by a new
  DescrptDPA4.supports_charge_spin() capability method, mirrored
  exactly on supports_native_spin: DPAtomicModel caches
  self.supports_charge_spin and forward_atomic_graph only forwards the
  kwarg to descriptors that declare it, so dpa1/dpa2/dpa3's call_graph
  (which has no charge_spin parameter) never TypeErrors.
- uses_graph_lower() rejected charge_spin_embedding and
  bridging_switch outright; both branches are removed, so only the
  explicit disable_graph_lower() escape hatch gates the graph route
  off now. Bridging itself needed no threading -- the trunk already
  reads self.bridging_switch directly.
- make_model.call_common (dpmodel, shared by pt_expt) forced any
  charge_spin request onto the dense lower via
  "if cs is not None: graph_method = None"; removed, and charge_spin
  is threaded through _call_common_graph -> call_lower_graph on both
  the dpmodel and pt_expt overrides.

has_message_passing_across_ranks() is untouched: bridging models still
fail multi-rank fast (a rank cannot observe a ghost owner's full
outgoing-edge set for the per-node freeze fold).

Tests (source/tests/common/dpmodel/test_dpa4_call_graph.py):
- test_uses_graph_lower_feature_gates flipped to the new gate set.
- charge_spin sensitivity + graph-vs-dense parity at 1e-12 (nf=2,
  distinct per-frame conditioning -- pins the nf threading, not just
  presence).
- bridging frozen-sphere invariance + gate-disabled ablation, mirroring
  pt's TestSourceFreezePropagationGate at the descriptor level.
- model-level: get_model(add_chg_spin_ebd=True) actually reaches
  DescrptDPA4.call_graph under neighbor_graph_method="dense" (spied,
  not just output-changed, since the old bug was a silent dense
  fallback rather than an error).

Author: Han Wang

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -406,8 +406,10 @@ 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

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -137,6 +137,14 @@ def __init__(
         self.supports_native_spin: bool = getattr(
             self.descriptor, "supports_native_spin", lambda: False
         )()
+        # Same capability-guard pattern as ``supports_native_spin`` above, for
+        # the frame-level ``charge_spin`` FiLM kwarg: only DPA4's
+        # ``call_graph`` declares it (see ``DescrptDPA4.supports_charge_spin``);
+        # other descriptors' ``call_graph`` would ``TypeError`` on an
+        # unconditional ``charge_spin=`` kwarg.
+        self.supports_charge_spin: bool = getattr(
+            self.descriptor, "supports_charge_spin", lambda: False
+        )()
         super().init_out_stat()
 
     def has_chg_spin_ebd(self) -> bool:
@@ -336,8 +344,11 @@ 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.
@@ -363,15 +374,20 @@ def forward_atomic_graph(
         # 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()
-        # See ``self.supports_native_spin`` in ``__init__``: only forward the
-        # ``spin`` keyword to descriptors whose ``call_graph`` declares it.
+        # See ``self.supports_native_spin``/``self.supports_charge_spin`` in
+        # ``__init__``: only forward the ``spin``/``charge_spin`` keyword to
+        # descriptors whose ``call_graph`` declares it.
         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,
             **spin_kwargs,
+            **charge_spin_kwargs,
         )
         fparam_node = None
         if fparam is not None:

deepmd/dpmodel/descriptor/dpa4.py

@@ -1676,6 +1676,7 @@ def call_graph(
         type_embedding: Array | None = None,
         comm_dict: dict[str, Array] | None = None,
         spin: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, None]:
         """Graph-native descriptor forward on the flat node axis.
 
@@ -1700,6 +1701,18 @@ def call_graph(
             None. Consumed by ``spin_embedding`` (l=0 magnitude into the type
             embedding, l=1 into the backbone and per-edge source features).
             Ghost-free graphs need only per-local-atom spin.
+        charge_spin
+            Frame-level charge/spin conditioning with shape ``(nf, 2)`` (or a
+            shape ``_canonicalize_charge_spin`` can broadcast to it), or
+            ``None``. This is the SAME per-descriptor canonicalization the
+            dense ``call`` adapter applies (default-fill from
+            ``default_chg_spin`` when configured, shape validation,
+            broadcast to ``nf``); ``call_graph`` is the one owner of that
+            step on the graph route. ``nf`` is recovered from
+            ``graph.n_node.shape[0]`` (a static shape, safe under
+            ``torch.export``); each frame's node block must therefore hold
+            exactly ``N // nf`` nodes, which single-rank carry-all graphs
+            built from a rectangular ``(nf, nloc)`` input always satisfy.
 
         Returns
         -------
@@ -1714,8 +1727,15 @@ def call_graph(
             exchange (raised by the per-block leaf).
         """
         n_nodes = atype.shape[0]
+        nf = graph.n_node.shape[0]
+        charge_spin = self._canonicalize_charge_spin(
+            charge_spin,
+            nf=nf,
+            dtype=graph.edge_vec.dtype,
+            device=array_api_compat.device(graph.edge_vec),
+        )
         x_scalar, _ = self._call_graph_common(
-            graph, atype, spin=spin, comm_dict=comm_dict
+            graph, atype, nf=nf, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict
         )
         # ``_call_graph_common`` returns the read-out with its SO(3) singleton
         # axes still attached, shape (n_nodes, 1, 1, channels); flatten to the
@@ -2332,19 +2352,18 @@ def uses_graph_lower(self) -> bool:
         Returns
         -------
         bool
-            False when the escape hatch has been pulled or when the model
-            uses conditioning inputs (charge/spin FiLM, SFPG bridging) that
-            still ride only the dense ``call`` signature. Native spin
-            (``spin_embedding``) IS supported on the graph lower: it is
-            threaded through ``call_graph`` like any other per-node input.
+            False only when the escape hatch has been pulled
+            (``disable_graph_lower()`` / ``_graph_lower_disabled``). Every
+            conditioning input DPA4 supports -- native spin
+            (``spin_embedding``), charge/spin FiLM (``charge_spin_embedding``),
+            and SFPG bridging (``bridging_switch``) -- rides the graph lower:
+            spin and charge_spin are threaded through ``call_graph`` like any
+            other per-node/per-frame input, and bridging is applied inside
+            the shared ``_call_graph_impl`` trunk with no extra threading (it
+            reads ``self.bridging_switch`` directly). Bridging models still
+            fail multi-rank fast via ``has_message_passing_across_ranks``.
         """
-        if self._graph_lower_disabled:
-            return False
-        if self.charge_spin_embedding is not None:
-            return False
-        if self.bridging_switch is not None:
-            return False
-        return True
+        return not self._graph_lower_disabled
 
     def uses_compact_edge_pairs(self) -> bool:
         """DPA4 attention is a per-edge scatter softmax; no pair axis."""
@@ -2354,6 +2373,10 @@ def supports_native_spin(self) -> bool:
         """DPA4 accepts a per-node ``spin`` on ``call_graph`` (native magnetic conditioning); other descriptors do not."""
         return True
 
+    def supports_charge_spin(self) -> bool:
+        """DPA4 accepts a frame-level ``charge_spin`` on ``call_graph`` (FiLM conditioning); other descriptors do not."""
+        return True
+
     def disable_graph_lower(self) -> None:
         """Route this descriptor through the legacy dense lower."""
         self._graph_lower_disabled = True

deepmd/dpmodel/model/make_model.py

@@ -310,6 +310,16 @@ def call_common(
                 The coordinates correction for virial.
                 shape: nf x (nloc x 3)
 
+            charge_spin
+                Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or
+                ``None``. Both the dense (nlist) and NeighborGraph lowers
+                consume it (currently only DPA4/SeZM); the graph route no
+                longer forces this model onto dense (former ``cs -> dense``
+                gate removed) -- it threads through
+                ``_call_common_graph``/``call_lower_graph`` to the
+                descriptor's ``call_graph``, gated per-descriptor by
+                ``supports_charge_spin``.
+
             spin
                 Per-local-atom spin, ``(nf, nloc, 3)``, or ``None``. Only the
                 NeighborGraph lower consumes it (native magnetic conditioning,
@@ -380,10 +390,6 @@ def call_common(
                         "pass one or the other"
                     )
                 graph_method = None
-            # the graph lower does not consume charge_spin yet -> keep those
-            # models on dense (a None check, so it stays jit/export-safe)
-            if cs is not None:
-                graph_method = None
             # model-level spin rides ONLY the NeighborGraph lower
             if sp is not None and graph_method is None:
                 raise NotImplementedError(
@@ -402,6 +408,7 @@ def call_common(
                     graph_method,
                     do_atomic_virial,
                     spin=sp,
+                    charge_spin=cs,
                 )
             else:
                 # legacy dense-nlist path (builds the extended quartet)
@@ -467,6 +474,7 @@ def _call_common_graph(
             method: str,
             do_atomic_virial: bool = False,
             spin: Array | None = None,
+            charge_spin: Array | None = None,
         ) -> dict[str, Array]:
             """Carry-all graph forward (opt-in, Option B).
 
@@ -495,6 +503,12 @@ def _call_common_graph(
                 Per-local-atom spin, ``(nf, nloc, 3)``, or ``None``. Flattened
                 to the flat node axis ``(N, 3)`` and forwarded unchanged to
                 :meth:`call_lower_graph`.
+            charge_spin
+                Frame-level charge/spin conditioning, ``(nf, 2)`` or ``None``.
+                Unflattened (per-frame, not per-node) and forwarded unchanged
+                to :meth:`call_lower_graph`, whose ``n_node`` here is always
+                the rectangular ``full(nf, nloc)`` this method builds -- the
+                one shape the descriptor's per-frame FiLM division requires.
 
             Returns
             -------
@@ -549,6 +563,7 @@ def _call_common_graph(
                     else None
                 ),
                 spin=(xp.reshape(spin, (nf * nloc, 3)) if spin is not None else None),
+                charge_spin=charge_spin,
             )
             # Public ABI is rectangular (nf, nloc, *); the lower is flat
             # (N=nf*nloc, *).  Unravel per-atom keys here at the boundary.

deepmd/pt_expt/model/make_model.py

@@ -584,8 +584,11 @@ def forward_common_lower_graph(
                 multi-rank graphs the ghost rows are included; their values
                 are inert under the owned-node mask).
             charge_spin
-                charge/spin conditioning. Ignored in PR-A; accepted for ABI
-                stability with charge/spin-conditioned descriptors.
+                Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or
+                ``None``, forwarded to the atomic model's
+                ``forward_common_atomic_graph`` (and, from there, the
+                descriptor's ``call_graph`` for descriptors that declare
+                ``supports_charge_spin``; currently DPA4 only).
             spin
                 Per-node native spin, flat ``(N, 3)``, or ``None``. When given,
                 a SECOND autograd leaf is created next to ``edge_vec`` and
@@ -738,6 +741,7 @@ def _call_common_graph(
             method: str,
             do_atomic_virial: bool = False,
             spin: torch.Tensor | None = None,
+            charge_spin: torch.Tensor | None = None,
         ) -> dict[str, torch.Tensor]:
             """Carry-all graph forward with autograd force/virial (pt_expt override).
 
@@ -767,6 +771,10 @@ def _call_common_graph(
                 Flattened to ``(N, 3)`` and forwarded into
                 :meth:`forward_common_lower_graph`, completing the seam
                 ``call_common`` (dpmodel, shared) opens for the graph route.
+            charge_spin
+                Frame-level charge/spin FiLM conditioning, ``(nf, 2)`` or
+                ``None``. Unflattened (per-frame) and forwarded unchanged into
+                :meth:`forward_common_lower_graph`.
 
             Returns
             -------
@@ -818,6 +826,7 @@ def _call_common_graph(
                 fparam=fp,
                 aparam=ap_flat,
                 spin=spin_flat,
+                charge_spin=charge_spin,
             )
             # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom
             # outputs (N = nf * nloc for a carry-all rectangular graph).