feat(dpmodel): NeighborGraph support for dpa1 tebd_input_mode="strip"

deepmd/dpmodel/descriptor/dpa1.py

@@ -432,18 +432,24 @@ def uses_graph_lower(self) -> bool:
 
         The graph-native lower (``call_graph``) covers the factorizable path
         AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D)
-        with concat type-embedding and no type exclusion. Remaining ineligible
-        configs (``tebd_input_mode == "strip"``, ``exclude_types``) fall back
-        to the legacy dense path, so those models keep working unchanged.
+        with concat OR strip type-embedding. Remaining ineligible configs
+        (``exclude_types``, and compressed descriptors) fall back to the legacy
+        dense path, so those models keep working unchanged.
 
         Eligibility does NOT imply numerical interchangeability with the
         dense route for every config: with ``smooth_type_embedding=True``
         the carry-all graph attention is sel-independent by design and
         differs from the dense lower by up to ~1e-4 (see the Notes of
         :meth:`call_graph`).
         """
+        # compressed descriptors have no graph kernel (geo/tebd tabulation is
+        # dense-only); keep them on the legacy dense path.
+        if self.compress:
+            return False
+        # exclude_types stays dense (graph exclusion is owned elsewhere); strip is
+        # now graph-eligible (per-edge factorized embedding, no neighbor coupling).
         return (
-            self.se_atten.tebd_input_mode == "concat"
+            self.se_atten.tebd_input_mode in ("concat", "strip")
             and not self.se_atten.exclude_types
         )
 
@@ -575,7 +581,7 @@ def call(
         nall = xp.reshape(coord_ext, (nlist.shape[0], -1)).shape[1] // 3
         # graph-eligible configs route through the graph-native adapter (decision
         # #14: graph = single math source, dense call = thin adapter). Ineligible
-        # configs (attention, strip tebd, exclude_types) and the ghost case with
+        # configs (exclude_types, compressed descriptors) and the ghost case with
         # no mapping fall back to the legacy dense body. The graph needs `mapping`
         # to fold ghosts to local owners; without it only nall == nloc is valid.
         if self.uses_graph_lower() and (mapping is not None or nall == nloc):
@@ -680,8 +686,8 @@ def _call_dense(
         atype_ext: Array,
         nlist: Array,
     ) -> Array:
-        """Legacy dense descriptor body (the ineligible ``call`` path: attention,
-        strip tebd, exclude_types, or the no-mapping ghost case).
+        """Legacy dense descriptor body (the ineligible ``call`` path:
+        compressed descriptors, exclude_types, or the no-mapping ghost case).
 
         Parameters
         ----------
@@ -1749,17 +1755,17 @@ def call_graph(
         Notes
         -----
         Known limitations:
-        - ``tebd_input_mode == "concat"`` only (strip mode lands later);
+        - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense;
         - ``exclude_types`` is not yet supported and raises (lands in a later PR).
         """
         from deepmd.dpmodel.utils.neighbor_graph import (
             edge_env_mat,
             segment_sum,
         )
 
-        if self.tebd_input_mode not in ["concat"]:
+        if self.tebd_input_mode not in ["concat", "strip"]:
             raise NotImplementedError(
-                "graph path supports tebd_input_mode='concat' only (NeighborGraph PR-A)"
+                f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}"
             )
         if self.exclude_types:
             raise NotImplementedError(
@@ -1794,20 +1800,25 @@ def call_graph(
         )  # (E, 4), (E, 1) sw zeroed on padding
         # radial channel
         ss = rr[:, 0:1]  # (E, 1)
-        # neighbor / center type embeddings (concat mode); ghost type == owner type
-        # so gathering by the LOCAL owner (src) reproduces the dense neighbor tebd.
-        # NB: do NOT wrap in ``xp.asarray(..., device=dev)`` -- that DETACHES under
-        # torch and severs the type-embedding weight gradient (the tebd net would
-        # never train); type_embedding already lives on the model device.
-        tebd = type_embedding
-        atype_embd_nlist = xp.take(tebd, nei_type, axis=0)  # (E, tebd_dim)
-        if not self.type_one_side:
-            atype_embd_nnei = xp.take(tebd, center_type, axis=0)  # (E, tebd_dim)
-            ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1)
-        else:
-            ss = xp.concat([ss, atype_embd_nlist], axis=-1)
-        # embedding net (same weights as the dense path); applies on the last axis
-        gg = self.embeddings[0].call(ss)  # (E, ng)
+        if self.tebd_input_mode == "concat":
+            # neighbor / center type embeddings; ghost type == owner type so
+            # gathering by the LOCAL owner (src) reproduces the dense neighbor tebd.
+            # NB: do NOT wrap in ``xp.asarray(..., device=dev)`` -- that DETACHES
+            # under torch and severs the type-embedding weight gradient (the tebd
+            # net would never train); type_embedding already lives on the device.
+            tebd = type_embedding
+            atype_embd_nlist = xp.take(tebd, nei_type, axis=0)  # (E, tebd_dim)
+            if not self.type_one_side:
+                atype_embd_nnei = xp.take(tebd, center_type, axis=0)  # (E, tebd_dim)
+                ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1)
+            else:
+                ss = xp.concat([ss, atype_embd_nlist], axis=-1)
+            # embedding net (same weights as the dense path); applies on last axis
+            gg = self.embeddings[0].call(ss)  # (E, ng)
+        else:  # strip: factorized gg_s*gg_t + gg_s (per-edge; no neighbor coupling)
+            gg = self._graph_edge_gg_strip(
+                ss, center_type, nei_type, type_embedding, sw_e
+            )
         # transformer attention over each center's edges — mirrors the dense
         # self.dpa1_attention(gg, nlist_mask, input_r, sw), which also runs on
         # the UNMASKED gg (padding rows are neutralized afterwards).
@@ -1835,6 +1846,78 @@ def call_graph(
         rot_mat = gr[:, :, 1:]
         return grrg, rot_mat
 
+    def _graph_edge_gg_strip(
+        self,
+        ss: Array,
+        center_type: Array,
+        nei_type: Array,
+        type_embedding: Array,
+        sw_e: Array,
+    ) -> Array:
+        """Per-edge stripped-tebd embedding, op-for-op vs the dense strip branch.
+
+        Mirrors the ``tebd_input_mode == "strip"`` block of :meth:`call`: the
+        geometric net runs on the radial channel only (``gg_s``), the stripped
+        type-embedding net produces a per-type(-pair) factor (``gg_t``,
+        optionally switch-smoothed), and the two combine as
+        ``gg_s * gg_t + gg_s``. The compression branches (geo/tebd) are NOT
+        reached on the graph route: :meth:`DescrptDPA1.uses_graph_lower`
+        excludes compressed descriptors, so this kernel assumes no compression.
+
+        Parameters
+        ----------
+        ss
+            (E, 1) per-edge radial channel (``rr[:, 0:1]``).
+        center_type
+            (E,) center (dst) LOCAL atom type of each edge.
+        nei_type
+            (E,) neighbor (src) LOCAL atom type of each edge.
+        type_embedding
+            (ntypes_with_padding, tebd_dim) type-embedding table.
+        sw_e
+            (E, 1) smooth switch, zeroed on padding edges.
+
+        Returns
+        -------
+        gg
+            (E, ng) per-edge embedding feeding the attention / segment_sum.
+        """
+        assert self.embeddings_strip is not None
+        xp = array_api_compat.array_namespace(ss)
+        nt = self.tebd_dim
+        ntypes_with_padding = type_embedding.shape[0]
+        # geometric net on the radial channel only (dense: gg_s = cal_g(ss_scalar))
+        gg_s = self.embeddings[0].call(ss)  # (E, ng)
+        if self.type_one_side:
+            # one-side strip table indexed by NEIGHBOR type only
+            tt_full = self.cal_g_strip(type_embedding, 0)  # (ntypes_pad, ng)
+            gg_t = xp.take(tt_full, nei_type, axis=0)  # (E, ng)
+        else:
+            # two-side type-pair table; row = center * ntypes_pad + nei
+            # (dense builds the same (ntypes_pad**2, 2*nt) table, nei-fastest).
+            type_embedding_nei = xp.tile(
+                xp.reshape(type_embedding, (1, ntypes_with_padding, nt)),
+                (ntypes_with_padding, 1, 1),
+            )
+            type_embedding_center = xp.tile(
+                xp.reshape(type_embedding, (ntypes_with_padding, 1, nt)),
+                (1, ntypes_with_padding, 1),
+            )
+            two_side_type_embedding = xp.reshape(
+                xp.concat([type_embedding_nei, type_embedding_center], axis=-1),
+                (-1, nt * 2),
+            )
+            tt_full = self.cal_g_strip(
+                two_side_type_embedding, 0
+            )  # (ntypes_pad**2, ng)
+            # int64 for torch take (take_along/take requires Long indices)
+            idx = xp.astype(center_type * ntypes_with_padding + nei_type, xp.int64)
+            gg_t = xp.take(tt_full, idx, axis=0)  # (E, ng)
+        if self.smooth:
+            # dense: gg_t = gg_t * sw (per-neighbor); sw_e is (E, 1), zeroed on padding
+            gg_t = gg_t * sw_e
+        return gg_s * gg_t + gg_s
+
     def _graph_attention(
         self,
         gg: Array,

source/tests/common/dpmodel/test_dpa1_call_graph_block.py

@@ -114,3 +114,93 @@ def test_exclude_types_raises(self) -> None:
             dd.se_atten.call_graph(
                 ng, self.atype.reshape(-1), type_embedding=dd.type_embedding.call()
             )
+
+
+class TestDpa1BlockCallGraphStrip:
+    """Bit-exact parity between the graph-native ``call_graph`` and the dense
+    ``call`` for ``tebd_input_mode='strip'``.
+
+    Strip mode factorizes the per-neighbor feature as ``gg = gg_s*gg_t + gg_s``
+    (radial embedding times type-pair strip embedding); it has no neighbor-axis
+    coupling, so the graph translation is edge-for-edge and must be bit-exact
+    with the dense block on the SAME neighbor list.
+    """
+
+    def setup_method(self) -> None:
+        rng = np.random.default_rng(3)
+        self.nloc = 4
+        self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5
+        self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64)
+
+    def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDPA1:
+        return DescrptDPA1(
+            rcut=4.0,
+            rcut_smth=0.5,
+            sel=[20],  # non-binding sel: carry-all graph == dense on real neighbors
+            ntypes=2,
+            attn_layer=attn_layer,
+            axis_neuron=2,
+            neuron=[6, 12],
+            tebd_input_mode="strip",
+            type_one_side=type_one_side,
+            smooth_type_embedding=smooth,
+        )
+
+    def _assert_parity(self, dd: DescrptDPA1, compact: bool) -> None:
+        (
+            ext_coord,
+            ext_atype,
+            mapping,
+            nlist,
+        ) = extend_input_and_build_neighbor_list(
+            self.coord,
+            self.atype,
+            dd.get_rcut(),
+            dd.get_sel(),
+            mixed_types=dd.mixed_types(),
+            box=None,
+        )
+        tebd = dd.type_embedding.call()
+        nf, nall = ext_atype.shape
+        atype_embd_ext = np.reshape(
+            np.take(tebd, np.reshape(ext_atype, (-1,)), axis=0),
+            (nf, nall, dd.tebd_dim),
+        )
+        dense_g, *_ = dd.se_atten.call(
+            nlist,
+            ext_coord,
+            ext_atype,
+            atype_embd_ext=atype_embd_ext,
+            mapping=None,
+            type_embedding=tebd,
+        )
+        ng = from_dense_quartet(ext_coord, nlist, mapping, compact=compact)
+        graph_g, _rot_mat = dd.se_atten.call_graph(
+            ng,
+            np.reshape(ext_atype, (-1,)),
+            type_embedding=tebd,
+        )
+        assert not np.any(np.isnan(graph_g))
+        np.testing.assert_allclose(
+            graph_g.reshape(dense_g.shape),
+            dense_g,
+            rtol=1e-12,
+            atol=1e-12,
+        )
+
+    @pytest.mark.parametrize(
+        "type_one_side", [False, True]
+    )  # two-side vs one-side strip table
+    @pytest.mark.parametrize("smooth", [False, True])  # gg_t switch-smoothing branch
+    def test_strip_attn0_equals_dense(self, type_one_side, smooth) -> None:
+        """attn_layer=0: no attention, so strip parity is bit-exact for both smooth values."""
+        dd = self._make(type_one_side, smooth, attn_layer=0)
+        self._assert_parity(dd, compact=True)
+
+    @pytest.mark.parametrize(
+        "type_one_side", [False, True]
+    )  # two-side vs one-side strip table
+    def test_strip_attn2_equals_dense(self, type_one_side) -> None:
+        """attn_layer=2, smooth=False: bit-exact (avoids by-design smooth softmax divergence)."""
+        dd = self._make(type_one_side, smooth=False, attn_layer=2)
+        self._assert_parity(dd, compact=False)

source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py

@@ -99,14 +99,13 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None:
     @pytest.mark.parametrize(
         "kwargs",
         [
-            {"tebd_input_mode": "strip"},  # strip tebd: graph unsupported -> dense
             {"exclude_types": [(0, 1)]},  # type exclusion: graph unsupported -> dense
         ],
     )
     def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None:
-        """attn_layer=0 configs the graph can't handle (strip tebd, exclude_types)
-        must report uses_graph_lower()=False and run the dense body without
-        raising (regression: Task-3 routing previously raised NotImplementedError).
+        """attn_layer=0 configs the graph can't handle (exclude_types) must report
+        uses_graph_lower()=False and run the dense body without raising (strip is
+        now graph-eligible; its routing + parity is covered by TestDpa1StripRouting).
         """
         dd = DescrptDPA1(
             rcut=4.0, rcut_smth=0.5, sel=[30], ntypes=2, attn_layer=0, **kwargs
@@ -200,3 +199,89 @@ def test_call_graph_returns_flat_node_axis(self) -> None:
         n = atype_local.shape[0]
         assert grrg.shape[0] == n and grrg.ndim == 2
         assert rot_mat.shape[0] == n and rot_mat.ndim == 3
+
+
+class TestDpa1StripRouting:
+    """Strip is now graph-eligible: ``uses_graph_lower()`` admits it, ``call``
+    routes through the graph adapter, and the adapter (``static_nnei``) is
+    bit-exact with the legacy ``_call_dense`` for EVERY strip config -- including
+    ``smooth_type_embedding=True`` at ``attn_layer>0`` (the adapter reproduces
+    the dense phantom-neighbor terms, unlike the direct carry-all ``call_graph``).
+    """
+
+    def setup_method(self) -> None:
+        rng = np.random.default_rng(7)
+        self.nloc = 4
+        self.coord = rng.normal(size=(1, self.nloc, 3)) * 1.5
+        self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64)
+
+    def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDPA1:
+        return DescrptDPA1(
+            rcut=4.0,
+            rcut_smth=0.5,
+            sel=[20],  # non-binding
+            ntypes=2,
+            attn_layer=attn_layer,
+            axis_neuron=2,
+            neuron=[6, 12],
+            tebd_input_mode="strip",
+            type_one_side=type_one_side,
+            smooth_type_embedding=smooth,
+            resnet_dt=False,
+        )
+
+    def test_uses_graph_lower_strip_gate(self) -> None:
+        """The gate admits non-compressed strip; excludes compressed and exclude_types."""
+        dd = self._make(type_one_side=False, smooth=True, attn_layer=2)
+        assert dd.uses_graph_lower() is True  # strip is now graph-eligible
+        # negative contract: compression keeps the descriptor on the dense path
+        dd.compress = True
+        assert dd.uses_graph_lower() is False
+        dd.compress = False
+        # negative contract: exclude_types still forces dense (separate feature,
+        # owned by the pair-exclude PR; strip does NOT change its eligibility)
+        dd_excl = DescrptDPA1(
+            rcut=4.0,
+            rcut_smth=0.5,
+            sel=[20],
+            ntypes=2,
+            attn_layer=2,
+            tebd_input_mode="strip",
+            exclude_types=[(0, 1)],
+            resnet_dt=False,
+        )
+        assert dd_excl.uses_graph_lower() is False
+
+    @pytest.mark.parametrize("type_one_side", [False, True])  # strip table branch
+    @pytest.mark.parametrize(
+        "smooth", [False, True]
+    )  # switch-smoothing + smooth attention
+    @pytest.mark.parametrize("attn_layer", [0, 2])  # no-attn and multi-layer attention
+    def test_call_strip_graph_equals_dense(
+        self, type_one_side, smooth, attn_layer
+    ) -> None:
+        """The routed ``call`` (graph adapter) is bit-exact with ``_call_dense``."""
+        dd = self._make(type_one_side, smooth, attn_layer)
+        assert dd.uses_graph_lower() is True  # precondition: call routes to graph
+        (
+            ext_coord,
+            ext_atype,
+            mapping,
+            nlist,
+        ) = extend_input_and_build_neighbor_list(
+            self.coord,
+            self.atype,
+            dd.get_rcut(),
+            dd.get_sel(),
+            mixed_types=dd.mixed_types(),
+            box=None,
+        )
+        routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping)
+        dense = dd._call_dense(ext_coord, ext_atype, nlist)
+        assert len(routed) == len(dense)
+        for r, d in zip(routed, dense, strict=True):
+            if r is None:
+                assert d is None
+                continue
+            assert not np.any(np.isnan(r))
+            np.testing.assert_allclose(r, d, rtol=1e-12, atol=1e-12)