Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 36 additions & 16 deletions deepmd/dpmodel/descriptor/dpa1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
----------
Expand Down
90 changes: 66 additions & 24 deletions source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
12 changes: 7 additions & 5 deletions source/tests/common/dpmodel/test_dpa1_graph_attention_parity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
7 changes: 6 additions & 1 deletion source/tests/pd/model/test_dpa1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
7 changes: 6 additions & 1 deletion source/tests/pt/model/test_dpa1.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
Loading