refactor(pt): reuse dpmodel EnvMatStat #5139
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughDevice-placement and dtype consistency improvements across dpmodel utilities via array_api_compat; PyTorch env_mat_stat now re-exports dpmodel implementations; a public alias was added to PairExcludeMask; region.normalize_coord uses a device- and dtype-matched scalar. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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. Comment |
There was a problem hiding this comment.
Pull request overview
This PR refactors the PyTorch implementation to reuse the dpmodel EnvMatStat classes, eliminating code duplication. Additionally, it improves array API compatibility by adding device parameters to array creation functions.
- Replaced PT-specific
EnvMatStatandEnvMatStatSeimplementations with imports from dpmodel - Added
build_type_exclude_maskmethod alias in PT'sPairExcludeMaskto match dpmodel interface - Enhanced array API compatibility by adding device parameters to array creation functions in dpmodel utilities
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| deepmd/pt/utils/env_mat_stat.py | Removed duplicate implementation; now imports EnvMatStat classes from dpmodel |
| deepmd/pt/utils/exclude_mask.py | Added build_type_exclude_mask method alias to align with dpmodel interface |
| deepmd/dpmodel/utils/nlist.py | Added device parameters to array creation functions (eye, ones, arange, asarray) and replaced .size with array_api_compat.size() |
| deepmd/dpmodel/utils/env_mat_stat.py | Added device parameters to array creation functions (zeros, ones, arange) and replaced .size with array_api_compat.size() |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5139 +/- ##
=======================================
Coverage 81.93% 81.93%
=======================================
Files 712 712
Lines 72864 72831 -33
Branches 3617 3616 -1
=======================================
- Hits 59700 59675 -25
+ Misses 12001 11993 -8
Partials 1163 1163 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/utils/exclude_mask.py (1)
115-148: Add device placement forself.type_maskand include missingaxis=0parameter.The device mismatch is real:
self.type_maskremains a host NumPy array whilenlist/atype_extmay be on GPU. This is evidenced by existing workarounds inenv_mat_stat.py(lines 196–198) that manually converttype_maskbefore callingbuild_type_exclude_mask. The fix should move this conversion into the method itself.Additionally, line 145 is missing the
axis=0parameter used inAtomExcludeMask.build_type_exclude_mask(line 57), which should be added for consistency.Proposed fix
type_ij = xp.reshape(type_ij, (nf, nloc * nnei)) + type_mask = xp.asarray( + self.type_mask, + device=array_api_compat.device(atype_ext), + ) mask = xp.reshape( - xp.take(self.type_mask[...], xp.reshape(type_ij, (-1,))), + xp.take(type_mask, xp.reshape(type_ij, (-1,)), axis=0), (nf, nloc, nnei), ) return mask
🤖 Fix all issues with AI agents
In @deepmd/dpmodel/utils/env_mat_stat.py:
- Around line 99-110: The ValueError raised when self.last_dim is invalid
contains a typo ("raial-only") and a vague message; update the exception in the
branch that sets radial_only to instead raise a clearer message like "last_dim
must be 1 (radial-only) or 4 (full descriptor)" to correct the typo and tighten
the wording, ensure the ValueError is a single concise string (no trailing
period needed), and keep the surrounding logic that sets radial_only based on
self.last_dim intact.
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/env_mat_stat.py (1)
193-207: Nice: reuse a single PairExcludeMask + movetype_maskonto the right device; consider sourcing device fromextended_atypeRight now you use
device=array_api_compat.device(atype); usingextended_atype(ornlist) would be a bit clearer/safer if any future refactor changes whereatypelives.Small refinement
@@ pair_exclude_mask.type_mask = xp.asarray( pair_exclude_mask.type_mask, - device=array_api_compat.device(atype), + device=array_api_compat.device(extended_atype), )
📜 Review details
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/dpmodel/utils/env_mat_stat.pydeepmd/dpmodel/utils/exclude_mask.py
🧰 Additional context used
🧠 Learnings (1)
📓 Common learnings
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-09-19T04:25:12.408Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4406
File: deepmd/dpmodel/array_api.py:51-53
Timestamp: 2024-11-23T00:01:06.984Z
Learning: In `deepmd/dpmodel/array_api.py`, the `__array_api_version__` attribute is guaranteed by the Array API standard to always be present, so error handling for its absence is not required.
Learnt from: njzjz
Repo: deepmodeling/deepmd-kit PR: 4204
File: deepmd/dpmodel/fitting/general_fitting.py:426-426
Timestamp: 2024-10-10T22:46:03.419Z
Learning: In `deepmd/dpmodel/fitting/general_fitting.py`, when using the Array API and `array_api_compat`, the `astype` method is not available as an array method. Instead, use `xp.astype()` from the array namespace for type casting.
🪛 Ruff (0.14.10)
deepmd/dpmodel/utils/env_mat_stat.py
104-106: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (3)
deepmd/dpmodel/utils/env_mat_stat.py (3)
58-65:array_api_compat.size(vv)is a solid portability fix
111-129: Device-awarezero_mean/one_stddevinit looks correct
184-191: Thedevice=parameter onxp.arange()is standard-compliant and working correctly.The Array API standard specifies
device=as an optional parameter on creation routines likearange(). The deepmd-kit codebase already uses this pattern extensively across multiple files (nlist.py, exclude_mask.py, region.py, etc.) with no reported issues. The patterndevice=array_api_compat.device(atype)correctly ensures the created array is placed on the same device as the reference array—NumPy safely ignores it (CPU-only), while torch/cupy respect it for proper device placement.
3 participants
Summary by CodeRabbit
Refactor
Bug Fixes
New Features
✏️ Tip: You can customize this high-level summary in your review settings.