refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)#5733
Open
wanghan-iapcm wants to merge 41 commits into
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 80.87% 80.76% -0.12%
==========================================
Files 1003 1004 +1
Lines 112492 112756 +264
Branches 4236 4248 +12
==========================================
+ Hits 90981 91067 +86
- Misses 19986 20158 +172
- Partials 1525 1531 +6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
58f771f to
9f2c31a
Compare
Comment on lines
+14
to
+16
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
Comment on lines
+21
to
+23
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
Comment on lines
+25
to
+27
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
d5b5032 to
9c96566
Compare
Comment on lines
+36
to
+38
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
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).
Add buildPairExcludeTable / applyPairExclusion (graph) / applyPairExclusionNlist (dense) in commonPT.h, structurally mirroring the Python transforms (apply_pair_exclusion, apply_pair_exclusion_nlist) with the same argument order and variable names (type_ij, keep). DeepPotPTExpt::init rebuilds the flat (ntypes+1)^2 keep table from the pair_exclude_types metadata; the seam applies it before every model call (graph and dense, LAMMPS and standalone) as an idempotent backstop to the exclusion already compiled into the .pt2.
gen_dpa1_pairexcl.py exports graph-route and dense-route DPA1(attn_layer=0) .pt2 models with model-level pair_exclude_types=[[0,1]] plus a no-exclusion baseline; references come from Python DeepEval of each model. test_deeppot_dpa1_pairexcl_ptexpt.cc validates both C++ ingestion routes (applyPairExclusion / applyPairExclusionNlist) against the Python references at 1e-10 (fp64), cross-checks graph==dense, and proves the exclusion is active by comparing against the empty-table baseline. Wired into test_cc_local.sh. 8/8 tests pass locally.
…e and forward_common_atomic_graph
The SpinModel backbone (dpa1 attn_layer=0) was being routed to the carry-all graph path by the default-flip (decision deepmodeling#17) introduced in the graph-pair-exclude branch. Virtual/placeholder types injected by get_spin_model double the atom density, making the system sel-binding; the carry-all graph keeps neighbors the capped dense nlist discards, so SpinModel.call_common diverged from call_common_lower (the dense lower used by pt_expt eager inference) by ~2e-6. Fix: pass neighbor_graph_method='legacy' when SpinModel.call_common invokes the backbone, forcing the dense-nlist path until spin-graph support is explicitly implemented. The graph .pt2 export path already has a fail-fast guard for spin (serialization.py:963). Adds a regression test pinning that: - SpinModel.call_common total energy equals backbone(legacy) on the same spin-doubled inputs (exact bit-identity). - The backbone in graph mode gives a DIFFERENT energy at this density, confirming the fixture exercises the diverging regime.
The pair-exclude branch removed exclude_types from DescrptDPA1.uses_graph_lower(), which previously kept spin backbones (they inject exclude_types) on the dense path. As a side effect the descriptor-level dispatch routed spin .pt2/.pte export through the graph kernel, whose scatter/atomic_add tripped a torch-inductor CPU codegen assertion (23 errors in test_deep_eval_spin.py). Add an explicit disable_graph_lower() knob on DescrptDPA1 and set it structurally in SpinModel.__init__ (covers get_spin_model and both dpmodel/pt_expt deserialize paths, so it survives serialize round trips). The flag is not serialized; re-derived at construction. The neighbor_graph_method=legacy kwarg in call_common is kept as belt-and-braces.
for more information, see https://pre-commit.ci
…s CodeQL/CodeRabbit
- RTD build failed: numpydoc's strict See Also parser rejected the C++-twin
cross-reference prose entries ('Error parsing See Also entry ...'). Move the
cross-refs to Notes sections (free-form reST) in apply_pair_exclusion and
apply_pair_exclusion_nlist.
- main.py: drop stale 'no type exclusion' from the graph-eligibility comment
(exclude_types is now graph-native; matches the docstring + ValueError).
- test_graph_atomic_parity: remove dead ds/ft/am chain (CodeQL unused 'am').
- test_neighbor_graph_builder: drop redundant local 'import unittest' (CodeQL).
Replace the NotImplementedError guard with a real angle remap: after edge compaction, remap angle_index onto the compacted edge axis via an exclusive prefix-sum over surviving edges and drop any angle whose constituent edges were excluded. angle_mask set without angle_index is rejected (nothing to remap). Covers the dpa3/se_t angle channel for the eager/dynamic-nedge path; the compiled/C++ path stays mask-only.
…n the atomic model Model-level pair_exclude_types is a canonical NeighborGraph BUILD transform (decision deepmodeling#18): fold it into edge_mask where the graph is constructed, so the graph lower / exported .pt2 consumes a pre-excluded graph and never re-applies it. Previously it was applied inside forward_common_atomic_graph, which the exported lower routes through -> the .pt2 double-applied it (C++ already masks at build via applyPairExclusion). Idempotent, so numerically identical; this removes the redundant table baked into the .pt2 and unifies all paths on the build-time seam. - dpmodel _call_common_graph: pass pair_excl to the builder (was relying on the atomic-model application; now aligned with pt_expt/C++) - pt_expt DeepEval _build_eval_graph: build pair_excl from the dpmodel/metadata and pass to all 4 backends (dense/ase/vesin/nv) -- previously .pt2 DeepEval inference relied on the in-model application, which is now gone - pt_expt compiled-training graph forward: pass pair_excl at build so eager==compiled holds for pair-excluded models - base_atomic_model.forward_common_atomic_graph: drop the apply_pair_exclusion call (+ unused import); the lower now consumes a pre-excluded graph - direct-lower callers (test_dpa1_graph_lower) apply exclusion at build via apply_pair_exclusion, mirroring the C++ transform after from_dense_quartet - new contract test: the lower does NOT re-apply model-level exclusion
…_pair_exclusion angle_index and angle_mask are a coupled pair. compact=True now validates them up front and raises ValueError on any inconsistent state (only one set; A dims disagree; angle_index not (2, A)) instead of silently defaulting a missing mask to all-real. Prevents a partial/mismatched pair from remapping silently wrong.
…ce-safe) The DeepEval graph builder applies model-level pair_exclude at build. Reusing the loaded dpmodel's pt_expt-wrapped pair_excl carries a torch (CUDA) type_mask buffer that fails to convert onto a numpy atype (dense/ase build path): 'can't convert cuda:0 tensor to numpy'. Build a fresh numpy PairExcludeMask from the exclude types instead -- it converts cleanly to numpy (dense/ase) or torch (vesin/nv) atype. Refresh the stale gen_dpa1_pairexcl docstring (exclusion is a build-time transform, not baked into the .pt2).
for more information, see https://pre-commit.ci
The debug probe confirmed the C++ graph-route pair exclusion IS active (edge_mask 30 -> 14 on the gtest system); the earlier gtest failure was a stale installed libdeepmd_cc.so, not a code bug. All 8 Dpa1PairExcl gtests pass on the Tesla T4 after a full rebuild + reinstall.
…eting A4 Decision deepmodeling#18/A4 (user, 2026-07-04): exclusion is applied ONCE when the neighbor list is built, not re-masked per forward — SAME design as the graph route. The A4 execution had wired pair_excl into the builders but left the atomic-model seam as an 'idempotent backstop' and never connected the callers, so in-tree dense ran on the backstop alone. This completes the port and removes the backstop: - base_atomic_model.forward_common_atomic: DROP the internal apply_pair_exclusion_nlist; the dense lower now consumes a pre-excluded nlist (contract documented; negative-contract test added) - dpmodel model_call_from_call_lower: pass pair_excl to builder.build (was never passed); extend_input_and_build_neighbor_list gains pair_excl - out-stat model_forward helper: build with pair_excl - SpinModel.process_spin_input_lower: the virtual-atom nlist extension IS the build site of the spin-extended nlist — fold the backbone's pair_excl in there (universal spin forward==forward_lower holds, 527 tests) - pt_expt: DeepEval dense builders (native strategy + inline + ASE) and the compiled-training dense branch build with pair_excl; _graph_pair_excl renamed _model_pair_excl (serves both routes) - jax: jax2tf TF wrapper gains a TF twin of the erasure (flat keep-table gather); jax2tf serialization + HLO wrapper (pair_excl rebuilt from model_def_script) + trainer prepare_input pass pair_excl at build - C++: applyPairExclusionNlist RESTORED as the single application site on the C++ dense route (its earlier removal was misaligned with A4); all 'idempotent backstop' comments rewritten as build-time ownership statements - tests: dense negative-contract test (lower must NOT re-apply); consistency TestEnerLower harness pre-excludes at build (legacy pt/pd re-apply internally = idempotent no-op, so cross-backend equality holds); test_dp_atomic_model excl-consistency feeds md0 a pre-excluded nlist; pt_expt graph-vs-dense lower parity pre-excludes the dense side
…orflow The jax2tf SavedModel wrapper previously carried a hand-written TensorFlow twin of the model-level pair-exclusion nlist transform (decision deepmodeling#18/A4), pinned to the canonical numpy/C++ implementation only by a value test. Replace it with a call to the canonical dpmodel apply_pair_exclusion_nlist through the vendored ndtensorflow array-API namespace -- the same mechanism the TF2 backend uses to run dpmodel code on TensorFlow. Unlike the neighbor-list *build* (which has data-dependent Python control flow and is deliberately kept as a TF twin, see jax2tf/nlist.py), the exclusion's only branch is on the static exclude_types config, so it traces cleanly under tf.saved_model.save. This removes the second implementation: the exclusion transform now has a single owner (dpmodel) reused across dpmodel / pt_expt / native-jax / TF2 / jax2tf, with the C++ ingestion seam the only remaining twin (unavoidable). Add source/tests/consistent/io/test_pair_exclude_savedmodel.py: exports a pair-excluded se_e2_a model to .savedmodel and checks energy vs the dpmodel reference (fp64), force/virial vs pytorch (the numpy dpmodel DeepEval does not compute usable forces), the identity/no-exclusion branch, and that the exclusion is genuinely active.
Replace the bespoke test_pair_exclude_savedmodel.py with an aligned
IOTest subclass (TestDeepPotPairExclude) in test_io.py: it supplies only
the model dict and inherits test_data_equal / test_deep_eval, which export
through every backend and cross-compare at rtol/atol 1e-12 (with the built-in
all-NaN skip covering the numpy dpmodel force path). Gated on
DP_TEST_TF2_ONLY, the mode in which test_deep_eval exercises the jax2tf
'.savedmodel' path -- the TF v1 backend raises NotImplementedError on
pair_exclude_types, so it must not run there.
That aligned cross-backend eval exposed a real gap: the tf2 backend
('.savedmodeltf') silently ignored model-level pair_exclude_types (its
SavedModel returned the non-excluded energy 2.847 vs the correct 2.843).
tf2 was outside the original PR scope (dpmodel/pt_expt/jax) and its outer
TF wrapper (deepmd/tf2/make_model.py) never folded the exclusion into the
nlist. Fix it the same way as every other backend: thread pair_excl into
model_call_from_call_lower and apply the canonical dpmodel
apply_pair_exclusion_nlist at the nlist-BUILD seam (the tensors are already
ndtensorflow arrays, so no wrap/unwrap). All five backends now agree.
…e-commit.ci autofix)
for more information, see https://pre-commit.ci
test_savedmodel_export_contains_xla_call_module passes a DummyModel with no atomic_model attribute; getattr(model.atomic_model, 'pair_excl', None) still evaluated model.atomic_model first and raised AttributeError. Guard the atomic_model access with a nested getattr in both the jax2tf and tf2 serialization callers.
CodeRabbit flagged the pt_expt graph-eligibility sentence as stale: it listed descriptor-level exclude_types as a disqualifier, but this PR removed that gate. uses_graph_lower() now only requires tebd_input_mode == 'concat' and folds exclude_types into the neighbor graph. Drop the exclude_types clause and state its support explicitly.
…(Piece A) The descriptor input stat (EnvMatStatSe.iter) applied model-level pair_exclude_types as an accumulation-DESELECT (excluded pairs dropped from the count), while the model forward feeds the descriptor a pre-excluded nlist (excluded pairs -> -1, treated like empty slots: env_mat 0, still counted). Those are different stat semantics. Align to the forward (decision deepmodeling#18/A4): fold model-level exclusion into the neighbor list EnvMatStatSe builds (pair_excl on extend_input_and_build_ neighbor_list) and drop the deselect block. Excluded pairs now zero-and-count, identical to descriptor-level exclude_types and to empty slots. This SHIFTS stored davg/dstd for models with pair_exclude_types (intended: it was misaligned with the forward). New invariant test asserts model-level == descriptor-level exclusion bit-identically, plus an exclusion-active control.
…ce B) The dpa1 forward computes its env matrix through the NeighborGraph (from_dense_quartet -> edge_env_mat); the input stat used the dense EnvMat. Route the dpa1 block's input stat through the SAME graph path so stat and forward share one env-matrix implementation, with both exclusions folded in exactly as the forward does (model-level via the pre-excluded nlist from Piece A; descriptor-level via the emask mask). BIT-IDENTICAL to the dense path, so stored davg/dstd are unchanged: from_dense_quartet(compact=False) reuses the same neighbor set + padding (row-major (frame,center,slot) edges), edge_env_mat mirrors EnvMat.call, and the (E,4) output reshapes 1:1 to the dense (nf,nloc,nsel,4) tensor. Opt-in via EnvMatStatSe(use_graph=True); se_e2_a/se_r are untouched. pt_expt inherits it via autowrap; legacy pt stays dense (bit-identical, so cross-backend parity holds). Tests: graph==dense stat bit-identical (1e-15) for se_e2_a and the dpa1 block, under no exclusion, model-level pair_exclude, and descriptor-level exclude.
_graph_env_mat (input stat) duplicated the dense-quartet -> (graph, atype_local) setup from DescrptDPA1._call_graph_adapter almost verbatim (coord reshape, mapping-None identity, from_dense_quartet compact=False, xp_take_first_n local atype). Extract it as neighbor_graph.graph_from_dense_ quartet(coord_ext, atype_ext, nlist, mapping) -> (graph, atype_local) and route both call sites through it. Pure extraction, bit-identical (dpa1 adapter + stat parity tests unchanged).
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
24cdad6 to
06e6ede
Compare
for more information, see https://pre-commit.ci
njzjz-bot
suggested changes
Jul 7, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Thanks for pushing this toward a build-seam implementation. I agree with the direction, but I don't think this is safe to merge yet: several lower/precomputed-entry paths still bypass the new pair-exclusion seam, and the input-stat cache can be reused across different model-level pair_exclude_types.
Requesting changes so the exclusion behavior is closed over all external entry points before merge. CI being green is expected here because the missing paths are mostly with-comm / direct edge / precomputed-neighbor cases that the current parity tests don't exercise.
Reviewed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5
| // never re-applies it; this is the single application site on the C++ | ||
| // dense route. | ||
| const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( | ||
| firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); |
Contributor
There was a problem hiding this comment.
This only covers the non-comm dense route. The use_with_comm dense branch above still calls run_model_with_comm(..., firstneigh_tensor, ...) directly, so a .pt2 with pair_exclude_types will still include excluded pairs in multi-rank / with-comm inference. Please apply applyPairExclusionNlist before the with-comm call too, and add a with-comm regression test.
| @@ -1238,13 +1268,26 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, | |||
| edge_tensors.edge_index_ext, edge_tensors.edge_mask, | |||
Contributor
There was a problem hiding this comment.
The lower_input_kind == "edge_vec" paths still pass the original edge_tensors.edge_mask into run_model_edges* without applying pair_exclude_table_. This affects both this non-comm path and the with-comm edge path above. Since the exported lower no longer re-applies model-level pair_exclude_types, the edge schema needs an applyPairExclusion-equivalent mask at this ingestion seam as well. Please be careful about the type space: fold_to_local=false needs extended atypes, while folded graph/edge inputs need local atypes.
| selection is a pure performance choice and results are unchanged. | ||
| """ | ||
| method = self._neighbor_graph_method | ||
| # Model-level ``pair_exclude_types`` is a graph-BUILD transform |
Contributor
There was a problem hiding this comment.
This wires model-level exclusion into the graph/nlist builders, but the direct edge fast path earlier in _prepare_lower_inputs still bypasses it: self._nlist_builder.build(..., return_mode="edges") returns an edge schema that is passed straight to the edge lower. vesin_neighbor_list.py also rejects pair_excl for return_mode="edges", so the caller needs to post-filter the returned edge schema / mask with the same graph-style exclusion. Please add a DeepEval lower_input_kind="edge_vec" + vesin/nv + pair_exclude_types test.
| # nlist. Excluded pairs then behave exactly like empty slots | ||
| # (env_mat 0, still counted) -- identical to descriptor-level | ||
| # exclude_types, replacing the previous accumulation-deselect. | ||
| pair_excl=pair_excl, |
Contributor
There was a problem hiding this comment.
Now that input statistics depend on model-level pair_exclude_types through the pre-excluded nlist, the stat cache key also needs to include a canonicalized exclusion set. EnvMatStatSe.get_hash() still hashes descriptor shape/cutoff/sel/etc. but not pair_exclude_types, so changing only the exclusion list can silently reuse stale stats.
Contributor
|
One more blocker I could not place as an inline review comment because the file is not touched in this PR diff:
Please either apply the model's pair exclusion to the converted Reviewed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5 |
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
3 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 transform (decision #18)Makes 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) refactored to call the transform for model-levelpair_exclude_types; stays as idempotent backstop.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 inDeepPotPTExpt::init. Exclusion is compiled into the exported graph (traced seam); the C++ call is an idempotent backstop. New gtest (8 tests) vs Python DeepEval reference at 1e-10.apply_pair_exclusionuseslogical_and+ bool cast (array_api_strict rejectedbool*bool), caught by the jax/strict consistency rows now traversing the graph path.Known limitations
pair_exclpath has no local oracle test (CUDA-only); to be validated on a GPU box.build_edge_exclude_maskstill returns int32 (bool cast at call sites; follow-up).🤖 Generated with Claude Code
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.