fix(jax): stabilize repflow dynamic selection export#5533
Merged
njzjz merged 4 commits intoJun 19, 2026
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a backend-internal ChangesStatic dynamic selection for RepFlows
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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/universal/dpmodel/descriptor/test_descriptor.py (1)
915-933: ⚡ Quick winCover the
use_loc_mapping=Falsestatic-index branch too.
_get_static_graph_index()changes its indexing stride whenuse_loc_mappingis disabled, but this helper always constructs the default mapped configuration. Extending this test to run both values would cover the second branch of the new static-dynamic layout logic that JAX now enables by default.♻️ Suggested test expansion
- def _make_dpa3(self, use_static_dynamic_sel: bool) -> DescrptDPA3: + def _make_dpa3( + self, + use_static_dynamic_sel: bool, + *, + use_loc_mapping: bool, + ) -> DescrptDPA3: # The switch is intentionally class-level and internal, so tests toggle # it only around construction and then restore the previous backend mode. old_use_static_dynamic_sel = DescrptBlockRepflows._use_static_dynamic_sel DescrptBlockRepflows._use_static_dynamic_sel = use_static_dynamic_sel try: return DescrptDPA3( **DescriptorParamDPA3( self.nt, self.rcut, self.rcut_smth, self.sel, ["O", "H"], smooth_edge_update=True, use_dynamic_sel=True, + use_loc_mapping=use_loc_mapping, ) ) finally: DescrptBlockRepflows._use_static_dynamic_sel = old_use_static_dynamic_sel def test_static_dynamic_sel_matches_packed_dynamic_sel(self) -> None: - packed = self._make_dpa3(False) - static = self._make_dpa3(True) + for use_loc_mapping in (True, False): + packed = self._make_dpa3(False, use_loc_mapping=use_loc_mapping) + static = self._make_dpa3(True, use_loc_mapping=use_loc_mapping) - packed_out = packed( - self.coord_ext, - self.atype_ext, - self.nlist, - mapping=self.mapping, - ) - static_out = static( - self.coord_ext, - self.atype_ext, - self.nlist, - mapping=self.mapping, - ) + packed_out = packed( + self.coord_ext, + self.atype_ext, + self.nlist, + mapping=self.mapping, + ) + static_out = static( + self.coord_ext, + self.atype_ext, + self.nlist, + mapping=self.mapping, + ) - np.testing.assert_allclose(packed_out[0], static_out[0], atol=self.atol) - np.testing.assert_allclose(packed_out[1], static_out[1], atol=self.atol) + np.testing.assert_allclose(packed_out[0], static_out[0], atol=self.atol) + np.testing.assert_allclose(packed_out[1], static_out[1], atol=self.atol) - valid_edge_mask = np.reshape(self.nlist != -1, (-1,)) - assert static_out[2].shape[0] == self.nf * self.nloc * sum(self.sel) - np.testing.assert_allclose( - packed_out[2], static_out[2][valid_edge_mask], atol=self.atol - ) - np.testing.assert_allclose( - packed_out[3], static_out[3][valid_edge_mask], atol=self.atol - ) - np.testing.assert_allclose( - packed_out[4], static_out[4][valid_edge_mask], atol=self.atol - ) + valid_edge_mask = np.reshape(self.nlist != -1, (-1,)) + assert static_out[2].shape[0] == self.nf * self.nloc * sum(self.sel) + np.testing.assert_allclose( + packed_out[2], static_out[2][valid_edge_mask], atol=self.atol + ) + np.testing.assert_allclose( + packed_out[3], static_out[3][valid_edge_mask], atol=self.atol + ) + np.testing.assert_allclose( + packed_out[4], static_out[4][valid_edge_mask], atol=self.atol + )Also applies to: 935-968
🤖 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/universal/dpmodel/descriptor/test_descriptor.py` around lines 915 - 933, The _make_dpa3 helper method currently only constructs the default mapped configuration, but _get_static_graph_index() has different indexing behavior when use_loc_mapping is disabled. Extend the _make_dpa3 method to accept a parameter for use_loc_mapping (similar to how it accepts use_static_dynamic_sel) and ensure it applies this parameter when constructing DescriptorParamDPA3. Then update all test methods that use _make_dpa3 (including those at lines 935-968) to run test assertions with both use_loc_mapping=True and use_loc_mapping=False so that both branches of the static-dynamic layout logic are covered.
🤖 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/universal/dpmodel/descriptor/test_descriptor.py`:
- Around line 915-933: The _make_dpa3 helper method currently only constructs
the default mapped configuration, but _get_static_graph_index() has different
indexing behavior when use_loc_mapping is disabled. Extend the _make_dpa3 method
to accept a parameter for use_loc_mapping (similar to how it accepts
use_static_dynamic_sel) and ensure it applies this parameter when constructing
DescriptorParamDPA3. Then update all test methods that use _make_dpa3 (including
those at lines 935-968) to run test assertions with both use_loc_mapping=True
and use_loc_mapping=False so that both branches of the static-dynamic layout
logic are covered.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: b5b5ca58-1ff0-4fbd-beea-b662619d0e7b
📒 Files selected for processing (4)
deepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/utils/network.pydeepmd/jax/descriptor/repflows.pysource/tests/universal/dpmodel/descriptor/test_descriptor.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5533 +/- ##
==========================================
- Coverage 82.18% 82.17% -0.01%
==========================================
Files 890 896 +6
Lines 101357 102688 +1331
Branches 4240 4343 +103
==========================================
+ Hits 83301 84385 +1084
- Misses 16754 16966 +212
- Partials 1302 1337 +35 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
approved these changes
Jun 16, 2026
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/dpmodel/descriptor/repflows.py`:
- Around line 280-285: In the warnings.warn() call that checks
self.use_dynamic_sel and self._use_static_dynamic_sel, add the stacklevel=2
parameter to the warnings.warn() function invocation. This parameter should be
added as a keyword argument after the warning message string to ensure the
warning points to the caller's instantiation site rather than the internal line
where the warning is issued.
🪄 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: 75695fc0-9c46-4dbf-ad41-383d69efbd80
📒 Files selected for processing (1)
deepmd/dpmodel/descriptor/repflows.py
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Jinzhe Zeng <njzjz@qq.com>
iProzd
approved these changes
Jun 18, 2026
iProzd
left a comment
iProzd
left a comment
Member
There was a problem hiding this comment.
Two non-blocking suggestions:
- The
a_sel**2angle materialization is the main caveat. Thewarnings.warnhelps, but it'd be good to also document in the export docs that JAX export materializesnf * nloc * a_sel**2angle slots, soa_selshould be kept modest for exportable models. - Patch coverage on
network.pyis low (~30%) — the newaggregatebranches (thenum_ownerconcat path and theaverage=Trueassert) aren't exercised. A small unit test targetingaggregatedirectly would prevent regressions.
LGTM otherwise.
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.
Summary
Validation
Summary by CodeRabbit
repflow/use_dynamic_selangle-selection constraints for JAX export and memory scaling.aggregateto compute bin counts only when required.aggregatewith explicitnum_owner.