diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index b338a88044..3659e30776 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -596,18 +596,25 @@ def call( sw The smooth switch function. """ - xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) - nloc = nlist.shape[1] - 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 (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): - return self._call_graph_adapter(coord_ext, atype_ext, nlist, mapping) - else: - return self._call_dense(coord_ext, atype_ext, nlist) + # The dense ``call`` always runs the legacy dense body -- it is the + # cross-backend consistency reference and must match the tf/pt/pd/jax + # dense descriptors bit-for-bit. It previously routed graph-eligible + # configs through ``_call_graph_adapter`` (decision #14), but the + # adapter is bit-exact ONLY in the trivial-statistics regime + # (``davg == 0``): the dense se_atten body leaks a phantom + # padding-neighbor ``-davg/dstd`` residual (``EnvMat.call`` subtracts + # ``davg`` AFTER the padding rows' geometry is weight-zeroed, and + # neither the empty ``exclude_types`` mask nor the attention layers + # re-mask it) that the graph path deliberately omits -- the graph + # output is the physically correct one, but ``call`` must reproduce + # the dense reference. The former gate also made ``mapping`` -- an + # argument that only enables ghost folding on graph routes -- silently + # change the numerics of a dense call. The graph-native route is + # reached exclusively through :meth:`call_graph` (pt_expt + # ``forward_atomic_graph`` and the graph ``.pt2``), never through + # ``call``. ``_call_graph_adapter`` is retained as the + # bit-exact-regime reference exercised by the adapter parity tests. + return self._call_dense(coord_ext, atype_ext, nlist) def _call_graph_adapter( self, @@ -616,14 +623,27 @@ def _call_graph_adapter( nlist: Array, mapping: Array | None, ) -> Array: - """Regime-1 dense->graph adapter (the eligible ``call`` path). + """Regime-1 dense->graph adapter. Builds a NeighborGraph from the dense quartet with the SHAPE-STATIC converter (``compact=False``, so this is jit/export-traceable -- no ``nonzero``), runs :meth:`call_graph`, and reconstructs the dense-shaped - ``sw``. Preserves the dense 5-tuple ABI exactly; masked invalid edges - contribute zero in ``call_graph``'s ``segment_sum`` so the output is - identical to the legacy dense body. + ``sw``. Preserves the dense 5-tuple ABI; masked invalid edges + contribute zero in ``call_graph``'s ``segment_sum``. + + Bit-exact vs :meth:`_call_dense` **only in the trivial-statistics + regime** (``davg == 0``). For nonzero ``davg`` the dense body leaks a + phantom padding-neighbor ``-davg/dstd`` residual into every padding + slot (``EnvMat.call`` subtracts ``davg`` AFTER the padding geometry is + weight-zeroed; with empty ``exclude_types`` nothing re-masks it, at + any ``attn_layer``) that the graph path deliberately omits. The graph + kernel additionally applies the slot-0 statistics ``[:, 0, :]`` to + every edge -- exact for real stat-computed tables (slot-uniform by + construction), not for artificially slot-varying ones. This is why + the dense :meth:`call` does NOT route here: it is the cross-backend + consistency reference. This method is retained as the + bit-exact-regime reference exercised by the adapter parity tests; the + production graph route is :meth:`call_graph`. Parameters ---------- diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index 0f072c70b3..5f0d7ceb03 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Full 5-tuple ABI parity between the graph-routed ``DescrptDPA1.call`` -(attn_layer=0, which now goes ``from_dense_quartet -> call_graph``) and the -legacy dense descriptor output captured BEFORE the swap, for binding AND -non-binding ``sel``. +"""Full 5-tuple ABI parity between ``DescrptDPA1._call_graph_adapter`` +(``from_dense_quartet -> call_graph``) and the legacy dense descriptor +output, for binding AND non-binding ``sel``. -The dense reference is reconstructed by calling the BLOCK directly -(``dd.se_atten.call``) and applying the descriptor-level ``concat_output_tebd`` -step by hand (mirroring dpa1.py), because ``dd.call`` itself now routes through -the graph for ``attn_layer == 0``. +The public ``dd.call`` does NOT route through the adapter (it is the +cross-backend dense reference; the adapter is bit-exact only in the +trivial-statistics regime these fresh models are in), so the adapter is +exercised directly. The dense reference is reconstructed by calling the +BLOCK directly (``dd.se_atten.call``) and applying the descriptor-level +``concat_output_tebd`` step by hand (mirroring dpa1.py). """ import numpy as np @@ -66,7 +67,7 @@ def _dense_reference(self, dd, ext_coord, ext_atype, nlist): @pytest.mark.parametrize("sel", [[30], [4]]) # non-binding AND binding def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: - """Graph-routed dd.call() returns the identical dense 5-tuple ABI.""" + """``_call_graph_adapter`` returns the identical dense 5-tuple ABI.""" dd = self._make(sel) ( ext_coord, @@ -81,10 +82,9 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: mixed_types=dd.mixed_types(), box=None, ) - # dense reference captured via the block (pre-swap behaviour) + # dense reference captured via the block ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - # the swapped public ABI: routes through the graph - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) assert len(out) == 5 # grrg np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) @@ -138,8 +138,8 @@ def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: ) # dense reference (calls block directly) ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - # graph-routed public call - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + # graph adapter, exercised directly (the public ``call`` is dense) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) assert len(out) == 5 np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) @@ -161,9 +161,10 @@ def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: """An eligible (concat) attn_layer=0 descriptor called with mapping=None - on a PERIODIC system (nall > nloc ghosts) must fall back to the dense - body and match it because the graph path requires an explicit ghost - mapping. + on a PERIODIC system (nall > nloc ghosts) must run the dense body and + match it (the public ``call`` is unconditionally dense; under the old + routing gate this pinned the mapping-less fallback, since the graph + path requires an explicit ghost mapping). """ dd = self._make([30]) box = np.eye(3, dtype=np.float64)[None] * 6.0 @@ -180,6 +181,44 @@ def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: out = dd.call(ext_coord, ext_atype, nlist, mapping=None) # must not IndexError np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) + def test_call_is_mapping_insensitive_with_nontrivial_stats(self) -> None: + """``call(..., mapping)`` == ``call(..., None)`` == ``_call_dense`` with + NONZERO ``davg`` and ghosts present. + + Regression for the removed graph-adapter routing gate: ``mapping`` only + enables ghost folding on graph routes and must never change the dense + numerics. Under the old gate, passing ``mapping`` silently switched an + eligible ``call`` to ``_call_graph_adapter``, whose output differs from + the dense reference whenever ``davg != 0`` (the dense phantom + padding-slot ``-davg/dstd`` leak the graph path deliberately omits) -- + so the same physical system gave different descriptors depending on + whether the caller supplied ``mapping``. Nonzero stats are injected + precisely because the fresh-model default (``davg == 0``) is the one + regime where the two routes coincide and the bug is invisible. + """ + rng = np.random.default_rng(7) + dd = self._make([30]) + nnei = sum(dd.get_sel()) + dd.se_atten.mean = rng.normal(size=(2, nnei, 4)) + dd.se_atten.stddev = 0.1 + np.abs(rng.normal(size=(2, nnei, 4))) + box = np.eye(3, dtype=np.float64)[None] * 6.0 + 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=box, + ) + assert ext_atype.shape[1] > self.nloc # ghosts present + with_mapping = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + without_mapping = dd.call(ext_coord, ext_atype, nlist, mapping=None) + dense = dd._call_dense(ext_coord, ext_atype, nlist) + np.testing.assert_allclose( + with_mapping[0], without_mapping[0], rtol=0.0, atol=0.0 + ) + np.testing.assert_allclose(with_mapping[0], dense[0], rtol=0.0, atol=0.0) + def test_single_rank_extension_keeps_type_invariant(self) -> None: """The ghost-free graph types a neighbor as ``atype[mapping[neighbor]]`` (its local owner). This is correct because a real single-rank extension @@ -188,7 +227,7 @@ def test_single_rank_extension_keeps_type_invariant(self) -> None: a periodic image of its owner and shares its type. This test pins that invariant (an inconsistent ``mapping`` like the universal fixture's old buggy permutation is NOT a valid single-rank extension) and confirms the - graph-routed ``call`` matches dense on the resulting quartet. + graph adapter matches dense on the resulting quartet. """ dd = self._make([30]) box = np.eye(3, dtype=np.float64)[None] * 6.0 @@ -208,7 +247,7 @@ def test_single_rank_extension_keeps_type_invariant(self) -> None: ext_atype[f], ext_atype[f][mapping[f]] ) # atype_ext[k] == atype[mapping[k]] ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + out = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) @@ -299,9 +338,12 @@ def test_uses_graph_lower_strip_gate(self) -> None: 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``.""" + """The strip graph adapter is bit-exact with ``_call_dense`` for a + fresh model (trivial statistics -- the adapter's bit-exact regime; + the public ``call`` itself is always dense). + """ dd = self._make(type_one_side, smooth, attn_layer) - assert dd.uses_graph_lower() is True # precondition: call routes to graph + assert dd.uses_graph_lower() is True # strip is graph-eligible ( ext_coord, ext_atype, @@ -315,10 +357,10 @@ def test_call_strip_graph_equals_dense( mixed_types=dd.mixed_types(), box=None, ) - routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + adapter = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) - assert len(routed) == len(dense) - for r, d in zip(routed, dense, strict=True): + assert len(adapter) == len(dense) + for r, d in zip(adapter, dense, strict=True): if r is None: assert d is None continue diff --git a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py index b156e0ff5b..439ba0fecf 100644 --- a/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py +++ b/source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py @@ -226,9 +226,11 @@ def test_se_atten_v2_is_graph_eligible(self) -> None: assert dd.uses_graph_lower() is True def test_se_atten_v2_graph_equals_dense(self) -> None: - """The graph-routed se_atten_v2 ``call`` is bit-exact with ``_call_dense`` + """The se_atten_v2 graph adapter is bit-exact with ``_call_dense`` (the ``static_nnei`` adapter reproduces the dense phantom terms despite - smooth=True) at a non-binding sel. + smooth=True) at a non-binding sel, for a fresh model (trivial + statistics -- the adapter's bit-exact regime; the public ``call`` + itself is always dense). """ from deepmd.dpmodel.descriptor.se_atten_v2 import ( DescrptSeAttenV2, @@ -243,10 +245,10 @@ def test_se_atten_v2_graph_equals_dense(self) -> None: ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( coord, atype, dd.get_rcut(), dd.get_sel(), mixed_types=True, box=None ) - routed = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + adapter = dd._call_graph_adapter(ext_coord, ext_atype, nlist, mapping) dense = dd._call_dense(ext_coord, ext_atype, nlist) - assert len(routed) == len(dense) - for r, d in zip(routed, dense, strict=True): + assert len(adapter) == len(dense) + for r, d in zip(adapter, dense, strict=True): if r is None: assert d is None continue diff --git a/source/tests/pd/model/test_dpa1.py b/source/tests/pd/model/test_dpa1.py index 285dd3d4cd..952f14c880 100644 --- a/source/tests/pd/model/test_dpa1.py +++ b/source/tests/pd/model/test_dpa1.py @@ -97,12 +97,17 @@ def test_consistency( atol=atol, err_msg=err_msg, ) - # dp impl + # dp impl. `mapping` is passed because that is the production + # invocation (DPAtomicModel.forward_atomic always forwards it); + # the dense `.call()` must give the same answer with or without + # it -- mapping only enables ghost folding on graph routes, it + # must never change the dense numerics. dd2 = DPDescrptDPA1.deserialize(dd0.serialize()) rd2, _, _, _, _ = dd2.call( self.coord_ext, self.atype_ext, self.nlist, + self.mapping, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(), diff --git a/source/tests/pt/model/test_dpa1.py b/source/tests/pt/model/test_dpa1.py index 36b7373bed..bdb306f4de 100644 --- a/source/tests/pt/model/test_dpa1.py +++ b/source/tests/pt/model/test_dpa1.py @@ -104,12 +104,17 @@ def test_consistency( atol=atol, err_msg=err_msg, ) - # dp impl + # dp impl. `mapping` is passed because that is the production + # invocation (DPAtomicModel.forward_atomic always forwards it); + # the dense `.call()` must give the same answer with or without + # it -- mapping only enables ghost folding on graph routes, it + # must never change the dense numerics. dd2 = DPDescrptDPA1.deserialize(dd0.serialize()) rd2, _, _, _, _ = dd2.call( self.coord_ext, self.atype_ext, self.nlist, + self.mapping, ) np.testing.assert_allclose( rd0.detach().cpu().numpy(),