refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)#5733
Conversation
|
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}
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.62% 79.52% -0.11%
==========================================
Files 1014 1015 +1
Lines 115533 115818 +285
Branches 4276 4291 +15
==========================================
+ Hits 91995 92102 +107
- Misses 21994 22164 +170
- Partials 1544 1552 +8 ☔ 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
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
wanghan-iapcm
force-pushed
the
feat-graph-pair-exclude
branch
from
d5b5032 to
9c96566
Compare
…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).
…xclude # Conflicts: # deepmd/dpmodel/descriptor/dpa1.py # deepmd/jax/model/hlo.py # source/tests/common/dpmodel/test_dpa1_call_graph_block.py # source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py # source/tests/pt_expt/descriptor/test_dpa1.py
…no recompile The pair-exclusion DPA3 MPI fixture was a second full inductor compile (~2 with- comm compiles, expensive). Model-level pair_exclude_types is applied at the C++ ingestion seam (from metadata.json) and the Python DeepEval build seam (from the serialized model.json), NOT baked into the exported graph -- so the compiled AOTI artifact (incl. the nested with-comm .pt2) is byte-identical to deeppot_dpa3_mpi.pt2. Derive the pairexcl fixture by copying that archive and patching pair_exclude_types into BOTH JSON blobs (metadata.json for C++, model.json for Python) so the two paths agree. Verified via DeepEval: the derived model is bit-identical to a full recompile (E=2.5311 vs baseline 2.4916) and differs from the no-exclusion baseline. deeppot_dpa3_mpi.pt2 stays the exact no-exclusion baseline for the reference and active-vs-baseline tests.
Apply the metadata-patch fixture derivation (a95a958) to gen_dpa1_pairexcl.py and factor the logic into gen_common.derive_pair_exclude_pt2 (used by both gen_dpa3.py and gen_dpa1_pairexcl.py). gen_dpa1_pairexcl.py generated three inductor-compiled models (_none/_graph/_nlist). _graph has the same weights and lower_kind as the _none graph baseline and differs only in pair_exclude_types, so it is now DERIVED from _none by patching the archive -- no third compile. _nlist is a different lower_kind (different exported graph) so it stays a compile. 3 compiles -> 2. Verified via DeepEval: the derived _graph is bit-identical to a full recompile (E active vs the _none baseline). gen_dpa3.py refactored to use the shared helper.
OutisLi
dismissed
their stale review
Retracted by author; posted prematurely and partly out of scope.
Audit of the PR's added docstrings against the package numpydoc convention
(CLAUDE.md): public functions were already compliant; fix the private/helper
idiom mismatches so they match their files' sibling methods.
- base_atomic_model._assert_{nlist,graph}_pair_excluded: add Parameters + Raises
sections (siblings like _finalize_atomic_ret use full numpydoc).
- env_mat_stat._graph_env_mat: add Parameters + Returns (siblings iter/__call__
use numpydoc).
- jax2tf/make_model.model_call_from_call_lower: revert docstring to its original
one-liner (the function documents no other params); the pair_excl rationale is
already an inline comment at the application site, so the single-param prose
was redundant + inconsistent.
- gen_common.derive_pair_exclude_pt2: concise summary + typed Parameters, move
the decision/verification cross-refs into a Notes section (not floating prose).
No Google-style/See-Also/RTD-breaking issues found.
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
The NumPy fail-safe guards, JAX-MD fix, C++ graph/dense with-comm coverage, one-time device-table upload, and associated non-vacuity tests are all verified at HEAD. However, the ingestion inventory is still incomplete: live/compiled TF2 paths do not pass pair_excl, and DeepPotJAX invokes the exported call_lower_* functions directly for InputNlist without applying the new build-time transform. Requesting changes until these two production paths are wired and covered. The branch also currently conflicts with master, so the pair-exclusion changes need to be retained when resolving the NeighborGraph-angle, HLO-property, and IO-test conflicts.
| charge_spin=charge_spin, | ||
| neighbor_list=neighbor_list, | ||
| ) | ||
| if pair_excl is not None: |
There was a problem hiding this comment.
[blocking] The new TF2 build-seam transform is not reached by the live model or trainer.
pair_excl defaults to None here, but deepmd/tf2/model/dp_model.py::call_common invokes this helper without passing model.atomic_model.pair_excl; the compiled trainer bypasses this helper and calls prepare_lower_inputs, whose signature and builders also have no pair_excl. Since the inherited dpmodel lower no longer reapplies exclusion and its fail-safe is NumPy-only, TF2 tensors silently retain excluded pairs in direct inference and both eager/compiled training. The current IO test only exercises tf2/utils/serialization.py, which does pass this argument, so it cannot catch the gap. Could we apply the canonical transform inside prepare_lower_inputs(..., pair_excl=...), thread it from both live call_common and the trainer, and add a live TF2/XLA non-vacuity test?
| # exclusion is a nlist-BUILD transform (decision #18/A4); | ||
| # the traced lower consumes a pre-excluded nlist. Guard | ||
| # atomic_model too: test doubles (DummyModel) lack it. | ||
| pair_excl=getattr( |
There was a problem hiding this comment.
[blocking] The JAX C++/LAMMPS InputNlist path bypasses this pre-exclusion.
This applies pair_excl only in the exported high-level call_* wrappers, while DeepPotJAX::compute(..., InputNlist, ...) invokes the SavedModel call_lower_* functions directly (source/api_cc/src/DeepPotJAX.cc:830-868). Those wrappers currently only format the nlist before entering the serialized JAX lower, which no longer reapplies model-level exclusion, and the NumPy-only guard is skipped during tracing. Consequently normal JAX evaluation is correct while the same model can silently include excluded pairs under JAX LAMMPS inference. Could we apply the canonical transform in the exported call_lower_* wrappers or at the C++ ingestion seam, and add a non-vacuous parity test between high-level call and raw call_lower/C++ InputNlist?
…xclude # Conflicts: # deepmd/dpmodel/utils/neighbor_graph/__init__.py # deepmd/jax/model/hlo.py # source/tests/consistent/io/test_io.py
The eager TF2 model (call_common in dp_model.py) built the nlist via model_call_from_call_lower but never passed pair_excl, so live/training inference silently included excluded type pairs (the SavedModel export already folds it in). Thread pair_excl through, matching the base dpmodel upper call and the SavedModel path (decision deepmodeling#18/A4).
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
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
I found two remaining issues in the build-seam refactor; details are inline.
| # ndtensorflow array, so no wrap/unwrap) via the canonical dpmodel | ||
| # helper, before the lower consumes it. Identity when nothing is | ||
| # excluded. | ||
| nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) |
There was a problem hiding this comment.
[blocking] TF2 compiled training still bypasses this transform. This fixes the live model.call_common path, but deepmd/tf2/train/trainer.py::_make_compiled_prepare_lower_batch calls prepare_lower_inputs() directly and then enters _call_common_lower_formatted. Since the lower no longer reapplies model-level exclusion, the actual compiled training path keeps excluded neighbors. A minimal type [0, 1] / exclusion (0, 1) case produces [[[1, -1], [0, -1]]] here, whereas the canonical transform produces all -1. Please move/thread pair_excl into prepare_lower_inputs, pass model.atomic_model.pair_excl from the trainer, and add an eager/XLA training non-vacuity regression test.
There was a problem hiding this comment.
Fixed. pair_excl is now threaded into prepare_lower_inputs as the single owner, and the trainer passes model.atomic_model.pair_excl at the compiled-training prepare seam:
- 9bdda5d —
prepare_lower_inputsfolds exclusion into the built nlist (the redundant application inmodel_call_from_call_loweris removed), andtrainer._make_compiled_prepare_lower_batchpassespair_excl. Includes the requested non-vacuity regression test (test_prepare_lower_inputs_folds_in_pair_exclusion): a type-[0,1] pair within cutoff yields a real neighbour without exclusion and is erased to −1 with it. - 4856eed — follow-up:
tf2/make_model.pymoved totf2/model/make_model.py(matching the other backends' layout), and the exclusion is applied at BUILD time inside the builder dispatch (custom strategies receivepair_exclviabuild(), the native inline builder appliesapply_pair_exclusion_nlistdirectly), mirroring the dpmodel/pt_expt convention instead of a post-hoc application.
Verified on a TF box (DP_TEST_TF2_ONLY=1): source/tests/tf2/test_training.py 30 passed; source/tests/consistent/model/test_ener.py 120 passed including the tf2 pair/atom-exclusion cases (which are no longer skipped as of c5b5345).
| builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() | ||
| extended_coord, extended_atype, nlist, mapping = builder.build( | ||
| cc, atype, bb, rcut, sel | ||
| cc, atype, bb, rcut, sel, pair_excl=pair_excl |
There was a problem hiding this comment.
[compatibility] This unconditionally breaks existing custom NeighborList implementations, even when the model has no exclusions. The previously published build(coord, atype, box, rcut, sel, return_mode=...) signature does not accept pair_excl; because this call always passes the keyword with None, such a strategy now fails with TypeError: ... got an unexpected keyword argument 'pair_excl' on every model. Please preserve the old interface by applying apply_pair_exclusion_nlist centrally to the returned quartet (preferred), or at least avoid passing the new keyword when it is None and provide a compatibility path for non-empty exclusion.
There was a problem hiding this comment.
Fixed, with a variation on the suggestion. Applying apply_pair_exclusion_nlist centrally on the returned quartet would put the transform outside the builder, inconsistent with the graph path where build_neighbor_graph(..., pair_excl=) owns exclusion — so the builder keeps ownership, and the keyword is passed only when set:
- 8988c05 / a947793 —
builder.build(cc, atype, bb, rcut, sel, **excl_kwargs)withexcl_kwargs = {} if pair_excl is None else {"pair_excl": pair_excl}.
Behavior split:
- no exclusion (every existing custom strategy): the old published
build(coord, atype, box, rcut, sel, return_mode=...)signature is called unchanged — noTypeError; - non-empty exclusion + legacy builder without the keyword: loud
TypeErrorinstead of silently including excluded pairs (fail-closed, which is the failure mode we want for this feature); - updated builders: exclusion applied at build, same as the graph path.
There was a problem hiding this comment.
Follow-up — the design settled differently from my previous reply: the conditional-kwargs indirection is dropped in a33ff79, and pair_excl=pair_excl is now passed to build() unconditionally and uniformly (dpmodel/tf2 make_model, pt_expt deep_eval, and the graph builders all use the same direct form).
Rationale: pair_excl is part of the NeighborList.build() contract — the base-class signature carries it with a None default. A custom strategy written against the pre-PR signature therefore fails loudly with TypeError rather than silently including excluded pairs; for a feature whose dangerous failure mode is silent inclusion, fail-loud is the behavior we want, and updating a custom build() to accept (or ignore via **kwargs) the new keyword is a one-line change on the implementer's side.
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
6b9b8ab to
8988c05
Compare
…_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.
| from deepmd.dpmodel.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) |
|
@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 |
…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.
5 participants
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.