fix(jax): fix DPA3 force NaN with edge_init_use_dist #4794
Merged
Conversation
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
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR replaces the direct use of xp.linalg.vector_norm with a safer variant to avoid NaN outputs when edge_init_use_dist is enabled.
- Swap
xp.linalg.vector_normforsafe_for_vector_normin the edge initialization path. - Ensures distance-based edge input computation handles zero or near-zero vectors without NaNs.
Comments suppressed due to low confidence (3)
deepmd/dpmodel/descriptor/repflows.py:529
- Consider adding unit tests for the
edge_init_use_distpath to verify thatsafe_for_vector_normcorrectly handles cases which previously resulted in NaN values.
edge_input = safe_for_vector_norm(diff, axis=-1, keepdims=True)
deepmd/dpmodel/descriptor/repflows.py:528
- [nitpick] Add a brief comment explaining why
safe_for_vector_normis used instead ofxp.linalg.vector_normto clarify the intention and avoid confusion.
if self.edge_init_use_dist:
deepmd/dpmodel/descriptor/repflows.py:529
- It looks like
safe_for_vector_normis not imported in this module. Add the necessary import or ensure it's defined in this scope.
edge_input = safe_for_vector_norm(diff, axis=-1, keepdims=True)
iProzd
approved these changes
Jun 10, 2025
Contributor
📝 WalkthroughWalkthroughThe update modifies the computation of vector norms in the Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant DescrptBlockRepflows
participant safe_for_vector_norm
User->>DescrptBlockRepflows: call(diff, edge_init_use_dist=True)
DescrptBlockRepflows->>safe_for_vector_norm: compute norm(diff, axis=-1, keepdims=True)
safe_for_vector_norm-->>DescrptBlockRepflows: return norm result
DescrptBlockRepflows-->>User: continue with edge_input initialization
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Pylint (3.3.7)deepmd/dpmodel/descriptor/repflows.pyNo files to lint: exiting. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🧰 Additional context used🧬 Code Graph Analysis (1)deepmd/dpmodel/descriptor/repflows.py (1)
⏰ Context from checks skipped due to timeout of 90000ms (28)
🔇 Additional comments (1)
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## devel #4794 +/- ##
==========================================
- Coverage 84.79% 84.79% -0.01%
==========================================
Files 698 698
Lines 67819 67819
Branches 3542 3541 -1
==========================================
- Hits 57509 57506 -3
Misses 9178 9178
- Partials 1132 1135 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Merged
via the queue into
deepmodeling:devel
with commit Jun 11, 2025
a6042c6
90 of 96 checks passed
ChiahsinChu
pushed a commit
to ChiahsinChu/deepmd-kit
that referenced
this pull request
Dec 17, 2025
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved the stability and reliability of distance calculations during edge initialization, enhancing overall robustness. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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 by CodeRabbit