test(pt_expt): graph attention make_fx (merge gate) + model force/virial parity

- test_make_fx_graph_attn: graph forward + autograd.grad at attn_layer=2
  traces under make_fx for BOTH smooth branches (the shape-static
  center_edge_pairs form is nonzero-free) — required since pt_expt compiled
  training routes eligible models through the graph lower.
- model-level graph-vs-legacy lower parity now parametrized over
  attn_layer {0, 2} (energy/force/virial/atom_virial, 1e-12 CPU).
- eligibility pins: attention+concat is graph-eligible; se_atten_v2
  (tebd_input_mode='strip') correctly stays dense (strip = later PR;
  the plan's 'se_atten_v2 inherits for free' did not hold).

Author: Han Wang

source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py

@@ -208,3 +208,20 @@ def test_torch_block_matches_numpy(self) -> None:
             static_nnei=nlist.shape[2],
         )
         np.testing.assert_allclose(out.numpy(), ref, rtol=1e-12, atol=1e-12)
+
+
+class TestGraphEligibility:
+    def test_attention_concat_is_graph_eligible(self) -> None:
+        assert _make(2).uses_graph_lower()
+
+    def test_strip_mode_stays_dense(self) -> None:
+        """se_atten_v2 (tebd_input_mode='strip') is NOT graph-eligible yet:
+        strip-mode graph support is a later PR; it must keep the dense route
+        (the PR-D plan's 'se_atten_v2 inherits for free' did not hold).
+        """
+        from deepmd.dpmodel.descriptor.se_atten_v2 import (
+            DescrptSeAttenV2,
+        )
+
+        dd = DescrptSeAttenV2(rcut=4.0, rcut_smth=0.5, sel=[20], ntypes=2, attn_layer=2)
+        assert not dd.uses_graph_lower()

source/tests/pt_expt/descriptor/test_dpa1.py

@@ -311,6 +311,65 @@ def fn(coord_ext, atype_ext, nlist, mapping):
             atol=atol,
         )
 
+    @pytest.mark.parametrize("smooth", [False, True])  # smooth attention branch
+    @pytest.mark.parametrize("prec", ["float64"])  # precision
+    def test_make_fx_graph_attn(self, prec, smooth) -> None:
+        """make_fx (export-readiness) of the GRAPH forward with attention.
+
+        MERGE BLOCKER (NeighborGraph PR-D): pt_expt compiled training routes
+        eligible models through the graph lower by default, so graph attention
+        (``attn_layer > 0``) must be fx-traceable — the shape-static
+        ``center_edge_pairs`` form keeps the pair enumeration ``nonzero``-free.
+        Covers both the smooth and non-smooth attention branches.
+        """
+        rng = np.random.default_rng(GLOBAL_SEED)
+        _, _, nnei = self.nlist.shape
+        davg = rng.normal(size=(self.nt, nnei, 4))
+        dstd = 0.1 + np.abs(rng.normal(size=(self.nt, nnei, 4)))
+
+        dtype = PRECISION_DICT[prec]
+        rtol, atol = get_tols(prec)
+        dd0 = DescrptDPA1(
+            self.rcut,
+            self.rcut_smth,
+            self.sel_mix,
+            self.nt,
+            attn_layer=2,
+            attn_dotr=True,
+            smooth_type_embedding=smooth,
+            precision=prec,
+            seed=GLOBAL_SEED,
+        ).to(self.device)
+        dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device)
+        dd0.se_atten.stddev = torch.tensor(dstd, dtype=dtype, device=self.device)
+        dd0 = dd0.eval()
+        coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device)
+        atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device)
+        nlist = torch.tensor(self.nlist, dtype=int, device=self.device)
+        mapping = torch.tensor(self.mapping, dtype=int, device=self.device)
+
+        def fn(coord_ext, atype_ext, nlist, mapping):
+            coord_ext = coord_ext.detach().requires_grad_(True)
+            rd = dd0(coord_ext, atype_ext, nlist, mapping)[0]
+            grad = torch.autograd.grad(rd.sum(), coord_ext, create_graph=False)[0]
+            return rd, grad
+
+        rd_eager, grad_eager = fn(coord_ext, atype_ext, nlist, mapping)
+        traced = make_fx(fn)(coord_ext, atype_ext, nlist, mapping)
+        rd_traced, grad_traced = traced(coord_ext, atype_ext, nlist, mapping)
+        np.testing.assert_allclose(
+            rd_eager.detach().cpu().numpy(),
+            rd_traced.detach().cpu().numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
+        np.testing.assert_allclose(
+            grad_eager.detach().cpu().numpy(),
+            grad_traced.detach().cpu().numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
+
     @pytest.mark.parametrize("shared_level", [0, 1])  # sharing level
     def test_share_params(self, shared_level) -> None:
         """share_params level 0: share all; level 1: share type_embedding only."""

source/tests/pt_expt/model/test_dpa1_graph_lower.py

@@ -91,7 +91,7 @@ def setup_method(self) -> None:
             [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device
         )
 
-    def _make_model(self) -> EnergyModel:
+    def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel:
         ds = DescrptDPA1(
             self.rcut,
             self.rcut_smth,
@@ -100,9 +100,13 @@ def _make_model(self) -> EnergyModel:
             neuron=[3, 6],
             axis_neuron=2,
             attn=4,
-            attn_layer=0,  # graph lower only supports attn_layer == 0
+            attn_layer=attn_layer,
             attn_dotr=True,
             attn_mask=False,
+            # smooth attention keeps sel-padding in the dense softmax
+            # denominator; the carry-all graph drops it BY DESIGN (PR-D), so
+            # exact graph-vs-dense parity requires smooth=False here.
+            smooth_type_embedding=smooth,
             activation_function="tanh",
             set_davg_zero=False,
             type_one_side=True,
@@ -165,13 +169,16 @@ def _prepare_lower_inputs(self, periodic: bool):
         mapping_t = torch.tensor(mapping, dtype=torch.int64, device=self.device)
         return ext_coord, ext_atype, nlist_t, mapping_t
 
+    @pytest.mark.parametrize("attn_layer", [0, 2])  # factorizable AND attention
     @pytest.mark.parametrize("periodic", [True, False])  # PBC vs non-PBC
     @pytest.mark.parametrize("do_av", [False, True])  # atom-virial off / on
-    def test_force_virial_parity_vs_legacy(self, periodic, do_av) -> None:
+    def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> None:
         """Graph lower energy/force/virial/atom_virial == legacy dense lower on
         the SAME neighbor set (regime-1 graph from from_dense_quartet).
+        attn_layer=2 exercises graph attention through model-level autograd
+        (smooth=False: exact carry-all parity regime, NeighborGraph PR-D).
         """
-        model = self._make_model()
+        model = self._make_model(attn_layer=attn_layer)
         model.eval()
         tol = (
             {"rtol": 1e-12, "atol": 1e-12}