refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)#5733
Merged
iProzd merged 66 commits intoJul 12, 2026
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds pair-exclusion support across neighbor-list and neighbor-graph paths, extends DPA1 graph-native attention and export/tracing behavior, updates spin routing, and applies matching pair-exclusion metadata in the C++ pt_expt inference seam with new tests. ChangesPython pair exclusion and graph-native DPA1 attention
C++ pt_expt pair-exclusion ingestion seam
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant DescrptDPA1
participant DescrptBlockSeAtten
participant center_edge_pairs
participant segment_softmax
DescrptDPA1->>DescrptBlockSeAtten: call_graph(graph, atype, static_nnei)
DescrptBlockSeAtten->>DescrptBlockSeAtten: apply_pair_exclusion(graph, atype, pair_excl)
DescrptBlockSeAtten->>center_edge_pairs: center_edge_pairs(dst, edge_mask, static_nnei)
DescrptBlockSeAtten->>segment_softmax: segment_softmax(scores, query_edge, mask)
segment_softmax-->>DescrptBlockSeAtten: normalized attention weights
DescrptBlockSeAtten-->>DescrptDPA1: grrg, rot_mat
sequenceDiagram
participant DeepPotPTExptInit
participant DeepPotPTExptCompute
participant buildPairExcludeTable
participant applyPairExclusion
participant applyPairExclusionNlist
DeepPotPTExptInit->>buildPairExcludeTable: pair_exclude_types
DeepPotPTExptCompute->>applyPairExclusion: graph edge_index, edge_mask, atype
DeepPotPTExptCompute->>applyPairExclusionNlist: nlist, atype_ext
applyPairExclusion-->>DeepPotPTExptCompute: filtered edge_mask
applyPairExclusionNlist-->>DeepPotPTExptCompute: filtered nlist
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pt_expt/descriptor/test_dpa1.py (1)
117-147: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPass
mappingtotorch.export.exporthere
test_exportablestill goes through the dense fallback becauseTestCaseSingleFrameWithNlistsetsnloc=3andnall=4, while this export call passes nomapping. That means the newexclude_typescase only covers the legacy dense exclusion mask, not the graph-nativeapply_pair_exclusionpath. Addmappingto the exported inputs so the parametrization exercises the intended route.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt_expt/descriptor/test_dpa1.py` around lines 117 - 147, test_exportable is missing the graph mapping input, so it still exercises the dense fallback instead of the graph-native exclusion path. Update the export setup in test_exportable to pass mapping into torch.export.export alongside dd0 and the existing inputs, using the test fixture’s mapping source so the exclude_types parametrization covers apply_pair_exclusion as intended.
🧹 Nitpick comments (7)
deepmd/dpmodel/model/make_model.py (1)
316-322: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring update looks accurate; stale example nearby.
Matches the new DPA1 graph-native attention behavior (attention layers now included in graph eligibility). Note the unchanged
_call_common_graphexception message a few dozen lines below ("e.g. dpa1 attn_layer=0") is now a narrower example than what this docstring describes — consider updating that message text for consistency in a follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/model/make_model.py` around lines 316 - 322, The exception message in _call_common_graph is now too narrow compared with the updated graph-native attention behavior described in the nearby docstring. Update the message text in _call_common_graph so it reflects the broader DPA1 attention-layer graph eligibility instead of only referencing the old “e.g. dpa1 attn_layer=0” example, keeping the wording consistent with the behavior documented in make_model.py.source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py (1)
166-179: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest comment overstates what is actually verified.
The comment says this block will "verify excluded pairs contribute sw == 0" and "check ... call_graph sw channel", but
call_graphonly returns(grrg, rot_mat)— it has noswoutput — and the actual assertions only check for NaN/Inf, not the claimed masking behavior. The earlierout[4]vsref[4]comparison (lines 162-164) already indirectly validates exclusion parity forswvia the dense reference, so this block is largely redundant and its comment is misleading about intent/coverage. Either remove the stale comment or replace it with an assertion that actually validates zeroed contributions from excluded pairs (e.g., inspect the block'sedge_mask/sw_eviase_atten.call_graphdirectly).♻️ Suggested comment fix (minimal)
- if exclude_types: - # verify excluded pairs contribute sw == 0 in the dense reference - # (atype=[0,1,0,1] -> pairs (0,1) and (1,0) should be masked) - # sw shape: (nf, nloc, nnei, 1); just check the graph output is also 0 - # for excluded-pair edges by checking call_graph sw channel + if exclude_types: + # additional sanity check on the raw call_graph output (no sw + # channel here; exclusion parity for sw is already verified via + # out[4] vs ref[4] above). graph = from_dense_quartet(ext_coord, nlist, mapping, compact=False) atype_local = self.atype.reshape(-1) - grrg_g, rot_mat_g = dd.call_graph( + grrg_g, _rot_mat_g = dd.call_graph( graph, atype_local, type_embedding=dd.type_embedding.call() ) # no nan/inf in output with exclusions applied assert not np.any(np.isnan(grrg_g)) assert not np.any(np.isinf(grrg_g))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py` around lines 166 - 179, The comment in test_dpa1_call_graph_descriptor is misleading because this block does not verify excluded-pair sw masking; call_graph only returns grrg and rot_mat, and the current assertions only check NaN/Inf. Update the test by either removing/rephrasing the stale comment to match the actual coverage, or add a real assertion for zeroed excluded-pair contributions by checking the relevant sw/edge-mask path through se_atten.call_graph or the returned graph data. The earlier out[4] vs ref[4] comparison already covers sw parity, so keep this block focused on what it truly validates.source/tests/common/dpmodel/test_neighbor_graph_builder.py (1)
419-427: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRedundant
import unittest.
unittestis already imported at the top of this file; the local re-import inside theexceptblock is unnecessary.🧹 Proposed cleanup
`@classmethod` def setUpClass(cls) -> None: try: import ase # noqa: F401 except ImportError as e: - import unittest - raise unittest.SkipTest("ase not installed") from e🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_neighbor_graph_builder.py` around lines 419 - 427, Remove the redundant local import inside test_neighbor_graph_builder’s setUpClass method: the file already imports unittest, so keep the ImportError handling but drop the inner import and use the existing unittest.SkipTest reference when ase is missing.Source: Linters/SAST tools
deepmd/dpmodel/utils/neighbor_graph/graph.py (1)
192-194: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring overstates what
compact=Truereplaces.The parameter doc says
edge_index,edge_vec,angle_index,angle_maskare all replaced whencompact=True. In practiceangle_index/angle_maskare never touched by the compact branch — the function only reaches the compaction step after confirming both areNone(otherwise it raisesNotImplementedError). Listing them as "replaced" could mislead a future implementer extending angle-compaction support into thinking this path already handles it.📝 Suggested doc fix
graph - The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, - ``edge_index``, ``edge_vec``, ``angle_index``, ``angle_mask``) are - replaced. + The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, + ``edge_index`` and ``edge_vec``) are replaced. ``angle_index`` / + ``angle_mask`` are never touched — compaction is rejected outright + when either is present (see the ``compact`` behavior below).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/utils/neighbor_graph/graph.py` around lines 192 - 194, The docstring for the neighbor graph parameter overstates the effect of compact=True by implying that angle_index and angle_mask are also replaced. Update the documentation in graph.py to say compact mode only compacts edge_index and edge_vec (along with edge_mask), and make it clear that angle_index and angle_mask are not handled by this branch because the code path only proceeds when they are None.deepmd/dpmodel/utils/neighbor_graph/ase_builder.py (1)
154-163: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPin device explicitly when converting
atypeforapply_pair_exclusion.
xp = array_api_compat.array_namespace(coord)followed byxp.asarray(atype)doesn't pin a device, unlike the analogouspair_exclwiring innv_graph_builder.pyandvesin_graph_builder.py, which both usetorch.as_tensor(atype, device=<coord's device>). Ifatypeisn't already a tensor on the same device ascoord(e.g. a CPU/numpyatypepaired with a CUDAcoord),xp.asarraywill silently produce a CPU tensor, which will then device-mismatch againstgraph.edge_index/edge_maskinsideapply_pair_exclusion.🔧 Suggested fix
if pair_excl is not None: import array_api_compat xp = array_api_compat.array_namespace(coord) - atype_flat = xp.reshape(xp.asarray(atype), (-1,)) + dev = array_api_compat.device(coord) + atype_flat = xp.reshape(xp.asarray(atype, device=dev), (-1,)) graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) return graph🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/utils/neighbor_graph/ase_builder.py` around lines 154 - 163, The atype conversion in the ASE neighbor graph path is not explicitly pinned to coord’s device, so apply_pair_exclusion can receive tensors on the wrong device. Update the ase_builder flow that builds graph and handles pair_excl to convert atype the same way as the nv_graph_builder and vesin_graph_builder paths: derive the device from coord and create atype on that device before flattening and passing it into apply_pair_exclusion. This keeps the device consistent with graph.edge_index and graph.edge_mask.source/tests/common/dpmodel/test_graph_atomic_parity.py (1)
318-344: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused model scaffolding.
amis never referenced here, soDescrptDPA1,InvarFitting, andDPAtomicModelcan be removed from this test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_graph_atomic_parity.py` around lines 318 - 344, The test builds unused model scaffolding that is never referenced, so remove the dead setup from test_apply_pair_exclusion_idempotent. Eliminate the DescrptDPA1, InvarFitting, and DPAtomicModel construction (including the am variable) and keep only the inputs actually needed for extend_input_and_build_neighbor_list, from_dense_quartet, and apply_pair_exclusion. Make sure the test still covers both the empty and non-empty pair_exclude_types branches.Sources: Coding guidelines, Linters/SAST tools
source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc (1)
101-159: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftTest coverage gap: LAMMPS
InputNlistingestion route not exercised for pair-exclusion.
check_against_ref/allTYPED_TESTs call the 6-argdp.compute(ener, force, virial, coord, atype, box), which routes toDeepPotPTExpt's standalone (no-nlist,build_nlist-based)compute()overload. The LAMMPS-styleInputNlistoverload — the actual pair-style ingestion seam, which cachesedge_index_tensor/firstneigh_tensoratago==0and recomputes geometry viacompactEdgeTensorsevery step before callingapplyPairExclusion/applyPairExclusionNlist— is never invoked here. A bug isolated to that branch's node/edge tensor construction (e.g. themulti_rank ? nall_real : nlocnode-count selection feedingapplyPairExclusion) wouldn't be caught by this suite.Consider adding a case that drives the
InputNlistoverload (mirroring the pattern intest_deeppot_dpa1_graph_ptexpt.cc) withpair_exclude_typesset, so both C++ ingestion entry points are validated against the Python reference.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc` around lines 101 - 159, Add coverage for the LAMMPS-style InputNlist ingestion path in this pair-exclusion test, because the current check_against_ref and TYPED_TESTs only exercise the 6-arg DeepPot::compute route. Introduce a test that calls the InputNlist compute overload on DeepPotPTExpt, using pair_exclude_types and matching the pattern used in test_deeppot_dpa1_graph_ptexpt.cc, so the edge/node tensor caching and applyPairExclusion/applyPairExclusionNlist branch are validated against the Python reference.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt_expt/entrypoints/main.py`:
- Around line 571-576: Update the stale inline comment in main by removing the
“no type exclusion” restriction so it matches the current graph-eligibility
behavior. Keep the note aligned with `_model_uses_graph_lower` in training.py
and the nearby ValueError message: describe graph lower as opt-in for
graph-eligible models (dpa1 with concat tebd, attention layers, and supported
exclude_types) and preserve the rest of the fail-fast/per-atom-virial
explanation.
In `@doc/model/train-se-atten.md`:
- Around line 160-164: Update the pt_expt training doc sentence describing
graph-eligible descriptors so it no longer says descriptor-level exclude_types
disqualifies the carry-all neighbor-graph path. Use the surrounding
se_atten/neighbor_graph_method explanation to state that mixed-type descriptors
with tebd_input_mode "concat" and no descriptor-level compression remain
graph-eligible, while exclude_types is not a blocking condition anymore. Keep
the dense-vs-graph parity note tied to smooth_type_embedding and attn_layer, but
make the eligibility rule consistent with the current behavior exercised by
test_exclude_types_graph_eligible_and_parity and dd.uses_graph_lower().
---
Outside diff comments:
In `@source/tests/pt_expt/descriptor/test_dpa1.py`:
- Around line 117-147: test_exportable is missing the graph mapping input, so it
still exercises the dense fallback instead of the graph-native exclusion path.
Update the export setup in test_exportable to pass mapping into
torch.export.export alongside dd0 and the existing inputs, using the test
fixture’s mapping source so the exclude_types parametrization covers
apply_pair_exclusion as intended.
---
Nitpick comments:
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 316-322: The exception message in _call_common_graph is now too
narrow compared with the updated graph-native attention behavior described in
the nearby docstring. Update the message text in _call_common_graph so it
reflects the broader DPA1 attention-layer graph eligibility instead of only
referencing the old “e.g. dpa1 attn_layer=0” example, keeping the wording
consistent with the behavior documented in make_model.py.
In `@deepmd/dpmodel/utils/neighbor_graph/ase_builder.py`:
- Around line 154-163: The atype conversion in the ASE neighbor graph path is
not explicitly pinned to coord’s device, so apply_pair_exclusion can receive
tensors on the wrong device. Update the ase_builder flow that builds graph and
handles pair_excl to convert atype the same way as the nv_graph_builder and
vesin_graph_builder paths: derive the device from coord and create atype on that
device before flattening and passing it into apply_pair_exclusion. This keeps
the device consistent with graph.edge_index and graph.edge_mask.
In `@deepmd/dpmodel/utils/neighbor_graph/graph.py`:
- Around line 192-194: The docstring for the neighbor graph parameter overstates
the effect of compact=True by implying that angle_index and angle_mask are also
replaced. Update the documentation in graph.py to say compact mode only compacts
edge_index and edge_vec (along with edge_mask), and make it clear that
angle_index and angle_mask are not handled by this branch because the code path
only proceeds when they are None.
In `@source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc`:
- Around line 101-159: Add coverage for the LAMMPS-style InputNlist ingestion
path in this pair-exclusion test, because the current check_against_ref and
TYPED_TESTs only exercise the 6-arg DeepPot::compute route. Introduce a test
that calls the InputNlist compute overload on DeepPotPTExpt, using
pair_exclude_types and matching the pattern used in
test_deeppot_dpa1_graph_ptexpt.cc, so the edge/node tensor caching and
applyPairExclusion/applyPairExclusionNlist branch are validated against the
Python reference.
In `@source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py`:
- Around line 166-179: The comment in test_dpa1_call_graph_descriptor is
misleading because this block does not verify excluded-pair sw masking;
call_graph only returns grrg and rot_mat, and the current assertions only check
NaN/Inf. Update the test by either removing/rephrasing the stale comment to
match the actual coverage, or add a real assertion for zeroed excluded-pair
contributions by checking the relevant sw/edge-mask path through
se_atten.call_graph or the returned graph data. The earlier out[4] vs ref[4]
comparison already covers sw parity, so keep this block focused on what it truly
validates.
In `@source/tests/common/dpmodel/test_graph_atomic_parity.py`:
- Around line 318-344: The test builds unused model scaffolding that is never
referenced, so remove the dead setup from test_apply_pair_exclusion_idempotent.
Eliminate the DescrptDPA1, InvarFitting, and DPAtomicModel construction
(including the am variable) and keep only the inputs actually needed for
extend_input_and_build_neighbor_list, from_dense_quartet, and
apply_pair_exclusion. Make sure the test still covers both the empty and
non-empty pair_exclude_types branches.
In `@source/tests/common/dpmodel/test_neighbor_graph_builder.py`:
- Around line 419-427: Remove the redundant local import inside
test_neighbor_graph_builder’s setUpClass method: the file already imports
unittest, so keep the ImportError handling but drop the inner import and use the
existing unittest.SkipTest reference when ase is missing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f66e001c-837b-4b14-a6b8-84d867cc8a56
📒 Files selected for processing (48)
deepmd/dpmodel/array_api.pydeepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/default_neighbor_list.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/ase_builder.pydeepmd/dpmodel/utils/neighbor_graph/builder.pydeepmd/dpmodel/utils/neighbor_graph/env.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/dpmodel/utils/neighbor_graph/pairs.pydeepmd/dpmodel/utils/neighbor_graph/segment.pydeepmd/dpmodel/utils/neighbor_list.pydeepmd/dpmodel/utils/nlist.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/pt_expt/utils/serialization.pydeepmd/pt_expt/utils/vesin_graph_builder.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydoc/model/train-se-atten.mdsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.ccsource/install/test_cc_local.shsource/tests/common/dpmodel/test_apply_pair_exclusion.pysource/tests/common/dpmodel/test_apply_pair_exclusion_nlist.pysource/tests/common/dpmodel/test_center_edge_pairs.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_call_graph_descriptor.pysource/tests/common/dpmodel/test_dpa1_graph_attention_parity.pysource/tests/common/dpmodel/test_graph_atomic_parity.pysource/tests/common/dpmodel/test_neighbor_graph_builder.pysource/tests/common/dpmodel/test_segment_softmax.pysource/tests/common/dpmodel/test_spin_model_legacy_routing.pysource/tests/infer/gen_dpa1_pairexcl.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/infer/test_graph_deepeval.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_linear_model.pysource/tests/pt_expt/utils/test_graph_pt2_metadata.pysource/tests/pt_expt/utils/test_neighbor_list.pysource/tests/pt_expt/utils/test_vesin_graph_builder.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5733 +/- ##
==========================================
- Coverage 79.65% 79.55% -0.11%
==========================================
Files 1015 1018 +3
Lines 115740 116180 +440
Branches 4275 4305 +30
==========================================
+ Hits 92198 92424 +226
- Misses 22001 22204 +203
- Partials 1541 1552 +11 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
July 5, 2026 14:52
58f771f to
9f2c31a
Compare
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
July 6, 2026 09:51
d5b5032 to
9c96566
Compare
added 14 commits
July 7, 2026 06:49
…ct=True) when angle fields are present
…exclusion; drop eligibility gate
…layer and type_one_side Add @pytest.mark.parametrize for attn_layer in [0, 2] and type_one_side in [False, True] to test_exclude_types_graph_parity. Also adds the missing parity assertion (graph vs dense at rtol=atol=1e-12, non-binding sel). Uses smooth_type_embedding=False to avoid the known by-design softmax denominator divergence in the dense smooth path.
Descriptor-level exclude_types is now graph-eligible (fully supported via apply_pair_exclusion). Remove 'no exclude_types' from four docstrings/error messages that list graph eligibility conditions. The gate condition was removed in the NeighborGraph implementation; only tebd_input_mode='concat' restriction remains. - deepmd/pt_expt/entrypoints/main.py: freeze_model docstring (~502) + ValueError message (~589) - deepmd/dpmodel/model/make_model.py: forward docstring (~317) - deepmd/pt_expt/train/training.py: _model_uses_graph_lower docstring (~591)
build_neighbor_graph, build_neighbor_graph_ase, build_neighbor_graph_vesin, build_neighbor_graph_nv all gain optional keyword-only pair_excl=None and compact=False; default path = geometric search then apply_pair_exclusion. _call_common_graph in pt_expt make_model wires atomic_model.pair_excl to every builder call so model-level pair_exclude_types is applied at build time (the atomic-model seam backstop stays as idempotent identity). Oracle tests assert set-equality of the valid-edge set between builder(pair_excl=X) and builder() + separate apply_pair_exclusion(X), for dense (2 id + 3 oracle cases) and ase (2 cases); vesin gets 4 new tests (2 identity, 2 oracle, parametrized over periodic).
…n for array_api_strict compat
…_neighbor_list/strategies (A4) Extract the inline pair-exclusion from base_atomic_model.forward_common_atomic into apply_pair_exclusion_nlist(nlist, atype_ext, pair_excl) in nlist.py. The seam is refactored to call the named helper (idempotent backstop remains). Add pair_excl=None to: - build_neighbor_list (dpmodel, nlist.py) - DefaultNeighborList.build - VesinNeighborList.build (pt_expt) - NvNeighborList.build (pt; CUDA-only, API parity) - NeighborList base class signature 12 new unit tests covering: None/empty identity, excluded pairs -> -1, -1 slot preservation, ghost-atom types, idempotence, torch namespace smoke, build_neighbor_list oracle equivalence, DefaultNeighborList oracle, VesinNeighborList oracle. NvNeighborList CUDA-only (not validated locally).
added 5 commits
July 10, 2026 23:19
test_exportable exported without mapping, so exclude_types only covered the dense fallback, not the graph-native apply_pair_exclusion path. Export both routes (with and without mapping) so the exclusion is traced on the graph lower too (CodeRabbit review).
C++ DeepPotJAX (which consumes both jax and tf2 .savedmodel via LAMMPS InputNlist) passed the raw nlist to the exported call_lower_* with no exclusion, so message-passing/direct-nlist inference silently included excluded type pairs (fail-open). The exported call_lower_* consumes a pre-excluded nlist by design (decision deepmodeling#18/A4), so fold exclusion in at the ingestion seam: - export a flat get_pair_exclude_types getter from both jax2tf and tf2 serialization (models exported before it baked exclusion into the graph, so a missing getter => identity is correct); - DeepPotJAX reads it at init and builds the keep table once; - filter the LAMMPS nlist vector (torch-free twin of applyPairExclusionNlist) before tensor construction. Hoist buildPairExcludeTable to the torch-free common.h so the libtorch apply-helpers (commonPT.h) and the TF-C-API DeepPotJAX share one builder.
…-test it Move the DeepPotJAX inline nlist exclusion loop into a shared torch-free helper applyPairExcludeNlistVec (common.h), the plain-vector twin of the libtorch applyPairExclusionNlist (commonPT.h). Add test_pair_exclude_table.cc pinning the keep-table convention and the vector filter (empty table == identity, both-direction exclusion, empty-slot/virtual-type handling). The test is torch-free so it runs in every C++ build, giving the DeepPotJAX seam logic direct coverage without a TF build.
The eager tf2 backend (deepmodeling#5598) skipped tf2 whenever pair_exclude_types or atom_exclude_types was non-empty, which hid that dp_model.call_common dropped pair_excl on the live path. Now that the build seam threads exclusion through the eager upper path, un-skip tf2 so the existing end-to-end consistency test exercises exclusion and guards against the drop (jax was already un-skipped and covered). Applies to both TestEner (upper) and TestEnerLower (lower).
Add test_deeppot_sea_pairexcl_jax.cc: loads jax .savedmodel + tf2 .savedmodeltf pairexcl variants (identical weights to the baseline, exclusion injected by inject_pair_exclude.py in convert-models.sh) and drives the LAMMPS InputNlist path. Asserts (1) excluded energy differs from baseline (exclusion active, not dropped) and (2) the InputNlist/lower path agrees with the coord-only/upper path (excludes the same pairs). Guards the DeepPotJAX integration that the torch-free unit test cannot. Runtime-skips when the JAX backend isn't built, so it compiles everywhere and runs under the TF-enabled CI build.
njzjz-bot
reviewed
Jul 10, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
I found two remaining issues in the build-seam refactor; details are inline.
added 4 commits
July 11, 2026 09:48
The live-path pair_excl fix accessed self.atomic_model directly, which raises AttributeError on tf2 model test doubles that lack it (broke test_tf2_dp_model_call_common_uses_tf2_helper). Guard atomic_model too, matching the SavedModel-export pattern.
… not in build() Passing pair_excl=... to builder.build() unconditionally broke custom NeighborList strategies whose published build(coord, atype, box, rcut, sel, return_mode=...) signature has no pair_excl param (TypeError on every model). Apply apply_pair_exclusion_nlist centrally on the returned nlist instead, so custom strategies keep the published signature (njzjz-bot review).
…aining The compiled training path (trainer._make_compiled_prepare_lower_batch -> prepare_lower_inputs) fed the lower directly without applying model-level pair_exclude_types, so compiled/XLA training kept excluded neighbors. Thread pair_excl into prepare_lower_inputs as the single owner (removing the now redundant application in model_call_from_call_lower), and pass model.atomic_model.pair_excl from the trainer. Add a non-vacuity regression test that the prepared nlist erases excluded cross-type pairs (njzjz-bot review).
Rework 01784ee: applying apply_pair_exclusion_nlist centrally in make_model put the transform outside the builder, inconsistent with the graph path where build_neighbor_graph owns exclusion. Restore builder ownership and pass the pair_excl keyword only when non-None: models without exclusion keep the old published build() signature (custom strategies unaffected, njzjz-bot's compatibility case), while a legacy builder combined with a non-empty exclusion fails loudly with TypeError instead of silently including excluded pairs (fail-closed).
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
July 11, 2026 01:53
6b9b8ab to
8988c05
Compare
added 2 commits
July 11, 2026 09:55
…_excl Move deepmd/tf2/make_model.py to deepmd/tf2/model/make_model.py, matching the dpmodel/pt/pt_expt/pd layout (model/make_model.py); added at top level by the eager-tf2 PR. Update all importers. Also align the pair-exclusion convention with dpmodel/pt_expt make_model: the nlist BUILDER owns exclusion. The custom-strategy branch passes pair_excl into neighbor_list.build() (keyword only when set, so legacy strategies keep working without exclusion and fail loudly with it), and the native inline builder applies apply_pair_exclusion_nlist at build time, instead of a post-hoc application at the end of prepare_lower_inputs.
Collaborator
Author
|
@iProzd All three items are addressed at HEAD (4856eed). 1. TF2 live/compiled paths now pass
2. DeepPotJAX InputNlist path now applies the build-time transform.
3. Conflicts with master are resolved (15858b6): the pair-exclusion changes are retained; the NeighborGraph-angle Also addressed from njzjz-bot's latest round (see inline replies): the |
added 2 commits
July 11, 2026 10:25
…ilder _nlist_builder is always the in-repo VesinNeighborList (never user-injected), so the direct keyword is correct here; the conditional-kwargs idiom in dpmodel/tf2 make_model exists only for third-party NeighborList strategies injected via the public neighbor_list= parameter.
…e contract Drop the conditional **excl_kwargs indirection in dpmodel/tf2 make_model and pass pair_excl=pair_excl uniformly (matching pt_expt deep_eval and the graph builders). pair_excl is part of the NeighborList.build() contract (base-class signature); a custom strategy predating it fails loudly with TypeError instead of silently including excluded pairs.
njzjz-bot
approved these changes
Jul 11, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Reviewed the current PR head. The remaining custom NeighborList compatibility concern is acceptable given the expected usage, and I found no other blocking issues.
Coding agent: Codex
Codex version: codex-cli 0.144.0-alpha.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
iProzd
enabled auto-merge
July 11, 2026 15:06
wanghan-iapcm
pushed a commit
to wanghan-iapcm/deepmd-kit
that referenced
this pull request
Jul 16, 2026
Resolves the both-add conflict in test_dpa1_call_graph_descriptor.py with master's pair-exclude merge (deepmodeling#5733, 98c3ab0): keep both tests -- master's test_eligible_no_mapping_with_ghosts_falls_back (docstring adapted: the public call is now unconditionally dense on this branch) and this branch's test_call_is_mapping_insensitive_with_nontrivial_stats.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pair
exclude_typesas a canonical NeighborGraph transformMakes pair-type exclusion a single canonical transform applied once at the neighbor-list/graph build seam, and uses it to add descriptor-level
exclude_typessupport to the dpa1 graph path (removing that eligibility gate), consistently across dpmodel, pt_expt, jax, and the C++ inference path.What changed
apply_pair_exclusion(graph, atype, pair_excl, *, compact=False)indeepmd/dpmodel/utils/neighbor_graph/: ANDsPairExcludeMask.build_edge_exclude_maskintograph.edge_mask.compact=Falseis mask-only (shape-static, export/AOTI-safe);compact=Truedrops masked edges (eager-only; raises on angle-carrying graphs). Idempotent.base_atomic_model): model-levelpair_exclude_typesis applied at BUILD time, soforward_common_atomic{,_graph}no longer re-apply it — they consume a pre-excluded nlist/graph (single-owner, not a backstop). To keep that from being fail-open, an eager (numpy) fail-safe assertion (_assert_nlist_pair_excluded/_assert_graph_pair_excluded) rejects a non-excluded input at the seam with a clear error; it is a no-op undertorch.export/ jaxjitand in compiled production (the exported-.pt2/ C++ paths are covered by ingestion-site tests). See the ingestion-path inventory below.exclude_types: theNotImplementedErrorand theuses_graph_lower()exclude condition are removed; exclusion applied insideDescrptBlockSeAtten.call_graphbefore the segment sums. Graph-vs-dense parity at non-binding sel is exact (rtol=atol=1e-12, attn_layer 0 and 2, type_one_side both).build_neighbor_graphand the pt_expt graph builders (dense/ase/vesin/nv) gain optionalpair_excl/compactwith default post-search application;_call_common_graphpasses model-level excludes at build time; oracle set-equality tests per available builder.apply_pair_exclusion_nlist(nlist, atype_ext, pair_excl)extracted from the inline seam code;build_neighbor_list+ Vesin/Nv/Default strategies gainpair_excl;return_mode='edges'+pair_exclfails fast.buildPairExcludeTable/applyPairExclusion/applyPairExclusionNlistinsource/api_cc/include/commonPT.h, mirroring the Python transforms (same arg order/variable names, cross-referenced docs);pair_exclude_typesserialized into.pt2metadata.jsonand rebuilt once inDeepPotPTExpt::init(device-resident table, uploaded once — no per-step H2D). Exclusion is applied at the C++ ingestion seam — the single owner on every run path; the exported lower consumes a pre-excluded input and never re-applies it. New gtest (8 tests) vs Python DeepEval reference at 1e-10, plus multi-rank LAMMPS exclusion tests (below).apply_pair_exclusionuseslogical_and+ bool cast (array_api_strict rejectedbool*bool), caught by the jax/strict consistency rows now traversing the graph path.Ingestion-path inventory (exclusion coverage)
Because the seam is now the single owner (fail-open if skipped), here is every entry point that builds a nlist/graph reaching the exclusion-owning lower, and how each applies exclusion:
_call_common(dense)build_neighbor_list(pair_excl=)_call_common_graphbuild_neighbor_graph(pair_excl=)exclude_typesapply_pair_exclusionincall_graphapply_pair_exclusion_nlist/ vesinpair_excl_build_eval_graph(pair_excl=)(dense/ase/vesin/nv)build_neighbor_graph(pair_excl=)in stat/forwardbuild_neighbor_list(pair_excl=)inEnvMatStatSe.itercall_lowerapply_pair_exclusion_nlistat the seam (fixed here)test_dense_neighbor_applies_model_pair_exclusionapplyPairExclusionNlistapplyPairExclusionNlist(fixed here)applyPairExclusionapplyPairExclusionGuard:
forward_common_atomic{,_graph}additionally carry an eager fail-safe assertion (numpy only) that rejects a non-excluded input, so any future dpmodel/jax ingestion miss fails loudly instead of silently including excluded pairs.Known limitations
pair_exclpath has no local oracle test (CUDA-only); to be validated on a GPU box.pair_exclude_types— pre-existing (predates this PR) and tracked as a separate fix.build_edge_exclude_maskstill returns int32 (bool cast at call sites; follow-up).Spin routing
Spin models auto-inject
exclude_types(virtual/placeholder types) into their backbone descriptor; before this PR that condition accidentally kept spin on the dense path. With exclude_types now graph-eligible, spin backbones flipped onto the carry-all graph route, which (a) diverges from the sel-capped reference on sel-binding spin systems and (b) trips a torch-inductor scatter codegen assertion during spin.pt2export. Fixed explicitly:DescrptDPA1.disable_graph_lower()(not serialized; re-derived structurally) is set inSpinModel.__init__— the single choke point coveringget_spin_model,SpinModel.deserialize, and the pt_expt spin classes — plus a belt-and-bracesneighbor_graph_method="legacy"atSpinModel.call_common. Regression tests pin the routing and its serialize→deserialize survival; the full spin export suite (23) and spin checkpoint-interop suite (12) are green.Verification
Full pt_expt suite: 1196 passed / 39 skipped / 3 failed — the 3 failures (
test_dpa4_freeze_to_pt2,test_dpa4_deep_eval_*) are byte-identical on the base commit (pre-existing torch-inductor dpa4 export issue on this box, unrelated). dpmodel exclusion suites 69 passed; consistency dpa1 99 passed/63 skipped (incl. jax + array_api_strict exclude rows); C++Dpa1PairExclgtest 8/8.Summary by CodeRabbit
.pt2metadata and enforcement at inference ingestion.