fix(dpmodel): use graph forward for output statistics - #5946
Open
Stardust0831 wants to merge 1 commit into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughThe atomic model wrapper now routes graph-capable models through neighbor-graph inference while retaining dense inference for other models. Tests cover parameter propagation, statistics, bias updates, direct graph equivalence, and empty systems. ChangesGraph atomic forwarding
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant AtomicModelWrapper
participant build_neighbor_graph
participant forward_common_atomic_graph
AtomicModelWrapper->>build_neighbor_graph: construct neighbor graph
build_neighbor_graph-->>AtomicModelWrapper: graph and flattened atom data
AtomicModelWrapper->>forward_common_atomic_graph: forward graph and model parameters
forward_common_atomic_graph-->>AtomicModelWrapper: atomic outputs
AtomicModelWrapper->>AtomicModelWrapper: restore atom dimensions
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt_expt/model/test_dpa1_graph_lower.py (1)
203-211: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated dense-route regression guard. Both
test_output_stat_forward_uses_graph_lowerandtest_change_out_bias_uses_graph_lowerdefine an identicalfail_densefunction and apply the same twomonkeypatch.setattrcalls (get_selandextend_input_and_build_neighbor_list) to prove the dense route is never used. One shared helper removes the duplication and keeps both guards synchronized if the dense-route entry points change.
source/tests/pt_expt/model/test_dpa1_graph_lower.py#L203-L211: replace this block with a call to a shared helper (e.g.,self._assert_dense_route_unused(monkeypatch, atomic_model)).source/tests/pt_expt/model/test_dpa1_graph_lower.py#L270-L277: replace this block with the same shared helper call.♻️ Proposed helper extraction
+ def _assert_dense_route_unused(self, monkeypatch, atomic_model) -> None: + def fail_dense(*args, **kwargs): + raise AssertionError("the graph statistics route must not use dense sel") + + monkeypatch.setattr(atomic_model, "get_sel", fail_dense) + monkeypatch.setattr( + "deepmd.dpmodel.utils.nlist.extend_input_and_build_neighbor_list", + fail_dense, + )Then in each test:
- def fail_dense(*args, **kwargs): - raise AssertionError("the graph statistics route must not use dense sel") - - monkeypatch.setattr(atomic_model, "get_sel", fail_dense) - monkeypatch.setattr( - "deepmd.dpmodel.utils.nlist.extend_input_and_build_neighbor_list", - fail_dense, - ) + self._assert_dense_route_unused(monkeypatch, atomic_model)🤖 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/model/test_dpa1_graph_lower.py` around lines 203 - 211, Extract the duplicated dense-route guard into a shared helper, such as _assert_dense_route_unused, in source/tests/pt_expt/model/test_dpa1_graph_lower.py. Have it define the failure callback and apply both monkeypatches for get_sel and extend_input_and_build_neighbor_list; replace the duplicated blocks at lines 203-211 and 270-277 with calls to this helper, passing monkeypatch and atomic_model.
🤖 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.
Nitpick comments:
In `@source/tests/pt_expt/model/test_dpa1_graph_lower.py`:
- Around line 203-211: Extract the duplicated dense-route guard into a shared
helper, such as _assert_dense_route_unused, in
source/tests/pt_expt/model/test_dpa1_graph_lower.py. Have it define the failure
callback and apply both monkeypatches for get_sel and
extend_input_and_build_neighbor_list; replace the duplicated blocks at lines
203-211 and 270-277 with calls to this helper, passing monkeypatch and
atomic_model.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 19393888-4818-44c2-8db2-de69b9682d7f
📒 Files selected for processing (2)
deepmd/dpmodel/atomic_model/base_atomic_model.pysource/tests/pt_expt/model/test_dpa1_graph_lower.py
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
1 participant
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.
Summary
This PR routes model predictions used by output-bias statistics through the graph lower when the atomic model declares
uses_graph_lower(), while preserving the existing dense neighbor-list path for all other models.Motivation
The output-statistics wrapper currently always calls
extend_input_and_build_neighbor_list(). For graph-native models,seldescribes normalization rather than a desired fixed neighbor capacity, so constructing a dense(nframes, nloc, nsel)neighbor list can allocate a large amount of padding that the model does not need. The regular graph execution path already avoids this representation by carrying only the neighbors found within the cutoff.Changes
uses_graph_lower()capability to select the statistics forward path without checking descriptor or model names.NeighborGraphdirectly from coordinates, atom types, the periodic box, cutoff, and model-level pair exclusions.atypeandaparamonto the graph node axis while keepingfparamandcharge_spinframe-level.(nframes * nloc, ...)to(nframes, nloc, ...)before the existing statistics reduction.Compatibility
This change does not add or modify any public API or configuration option. Traditional dense models retain their existing execution path and output behavior. Empty systems also retain the dense path because the existing dense neighbor-list implementation explicitly supports zero local atoms.
Tests
sel.change-by-statisticregression test using self-consistent labels.ruff format --check,ruff check, andgit diff --checkon the changed files.Summary by CodeRabbit