refactor(jax): auto-convert dpmodel modules#5527
Merged
Merged
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
CodeQL found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
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 (2)
📝 WalkthroughWalkthroughThis PR centralizes JAX-side attribute conversion in ChangesJAX Attribute-Setting Consolidation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/jax/model/dp_zbl_model.py (1)
3-5: ⚡ Quick winIf this import is registration-only, make that side effect explicit.
After removing
DPZBLModel.__setattr__, this file depends on the ZBL atomic-model wrapper being imported beforeatomic_modelassignments hit the centralized conversion path. The current symbol import hides that dependency; mirroring the underscore-prefixed# noqa: F401side-effect import style used in the other JAX wrapper files would make the contract much clearer. Based on the import patterns in the other JAX wrapper files in this PR and the model-level wiring summary.🤖 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/jax/model/dp_zbl_model.py` around lines 3 - 5, This import is registration-only and should make that side effect explicit: change the statement that brings in DPZBLLinearEnergyAtomicModel to an underscore-prefixed, noop import and add a noqa comment (e.g., import DPZBLLinearEnergyAtomicModel as _DPZBLLinearEnergyAtomicModel # noqa: F401) so the file documents the registration side effect required by DPZBLModel.__setattr__ removal and suppresses unused-symbol warnings.
🤖 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/jax/common.py`:
- Around line 107-134: The registration routine can be raced: modify
_ensure_registrations to use a threading.Condition (or Lock+Condition) to
serialize the first registration and block concurrent callers until the
registration pass completes; specifically, protect _REGISTRATIONS_IN_PROGRESS
and _REGISTRATIONS_READY with a module-level Condition, then in
_ensure_registrations acquire the condition and if _REGISTRATIONS_READY return,
if _REGISTRATIONS_IN_PROGRESS call cond.wait() in a loop until
_REGISTRATIONS_READY or in_progress is false, otherwise set
_REGISTRATIONS_IN_PROGRESS = True and release the condition while performing the
import_module loop, then in a finally re-acquire the condition, set
_REGISTRATIONS_IN_PROGRESS = False, set _REGISTRATIONS_READY = True only on
successful registration, and call cond.notify_all() so waiting callers resume
and see the final registry state used by try_convert_module/_DPMODEL_TO_JAX.
---
Nitpick comments:
In `@deepmd/jax/model/dp_zbl_model.py`:
- Around line 3-5: This import is registration-only and should make that side
effect explicit: change the statement that brings in
DPZBLLinearEnergyAtomicModel to an underscore-prefixed, noop import and add a
noqa comment (e.g., import DPZBLLinearEnergyAtomicModel as
_DPZBLLinearEnergyAtomicModel # noqa: F401) so the file documents the
registration side effect required by DPZBLModel.__setattr__ removal and
suppresses unused-symbol warnings.
🪄 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: b701da7a-cf1f-4f7a-9929-0bd77e43855c
📒 Files selected for processing (22)
deepmd/jax/atomic_model/base_atomic_model.pydeepmd/jax/atomic_model/dp_atomic_model.pydeepmd/jax/atomic_model/linear_atomic_model.pydeepmd/jax/atomic_model/pairtab_atomic_model.pydeepmd/jax/common.pydeepmd/jax/descriptor/dpa1.pydeepmd/jax/descriptor/dpa2.pydeepmd/jax/descriptor/dpa3.pydeepmd/jax/descriptor/hybrid.pydeepmd/jax/descriptor/repflows.pydeepmd/jax/descriptor/repformers.pydeepmd/jax/descriptor/se_atten_v2.pydeepmd/jax/descriptor/se_e2_a.pydeepmd/jax/descriptor/se_e2_r.pydeepmd/jax/descriptor/se_t.pydeepmd/jax/descriptor/se_t_tebd.pydeepmd/jax/fitting/fitting.pydeepmd/jax/model/dp_model.pydeepmd/jax/model/dp_zbl_model.pydeepmd/jax/utils/exclude_mask.pydeepmd/jax/utils/network.pydeepmd/jax/utils/type_embed.py
💤 Files with no reviewable changes (2)
- deepmd/jax/atomic_model/base_atomic_model.py
- deepmd/jax/model/dp_model.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5527 +/- ##
==========================================
- Coverage 82.19% 82.18% -0.01%
==========================================
Files 891 890 -1
Lines 101599 101357 -242
Branches 4242 4242
==========================================
- Hits 83507 83299 -208
+ Misses 16789 16754 -35
- Partials 1303 1304 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
approved these changes
Jun 14, 2026
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:
Tests:
Summary by CodeRabbit