refactor(pt): reuse dpmodel EnvMatStat #5139

njzjz · 2026-01-08T11:37:59Z

Summary by CodeRabbit

Refactor
- Consolidated duplicate utility implementations by re-exporting shared components to reduce duplication and simplify maintenance.
Bug Fixes
- Improved device and dtype consistency for array operations to ensure tensors are allocated and computed on the correct device across backends.
- Ensured numeric operations use matching dtype/device for reliable behavior.
New Features
- Added a convenient public alias for mask-building functionality.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

for more information, see https://pre-commit.ci

coderabbitai · 2026-01-08T11:39:46Z

📝 Walkthrough

Walkthrough

Device-placement consistency added across dpmodel utilities via array_api_compat; PyTorch env_mat_stat now re-exports dpmodel implementations; a public alias was added to PairExcludeMask. Internal numeric dtype/device handling in region.normalize_coord was adjusted.

Changes

Cohort / File(s)	Summary
Device handling in dpmodel utilities `deepmd/dpmodel/utils/env_mat_stat.py`, `deepmd/dpmodel/utils/nlist.py`	Replaced direct size/creation calls with `array_api_compat.size(...)` and added explicit `device=array_api_compat.device(...)` to array operations (eye, ones, arange, outer, asarray, asarray, etc.). Ensures tensors (zero_mean, one_stddev, type_idx/aidx, ghost computations, concatenation/padding) are allocated on the correct device.
PyTorch module re-export `deepmd/pt/utils/env_mat_stat.py`	Removed local EnvMatStat/EnvMatStatSe implementations and now import/re-export them from `deepmd.dpmodel.utils.env_mat_stat`; module reduced to `__all__ = ["EnvMatStat", "EnvMatStatSe"]`.
PyTorch utility alias `deepmd/pt/utils/exclude_mask.py`	Added `build_type_exclude_mask = forward` alias on `PairExcludeMask` (public alias only).
Region dtype/device consistency `deepmd/dpmodel/utils/region.py`	Changed remainder operand to `xp.ones((), dtype=icoord.dtype, device=array_api_compat.device(icoord))` to match dtype and device of `icoord` during `normalize_coord`.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

iProzd

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'refactor(pt): reuse dpmodel EnvMatStat' directly describes the primary change: refactoring the pt module to reuse EnvMatStat from dpmodel. This is clearly supported by the deepmd/pt/utils/env_mat_stat.py changes, which replace local implementations with imports from deepmd.dpmodel.utils.env_mat_stat.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

📝 Generate docstrings

📜 Recent review details

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between b652f03 and c690b5d.

📒 Files selected for processing (1)

deepmd/dpmodel/utils/region.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-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: 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: 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.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (40)

GitHub Check: Build wheels for cp311-win_amd64
GitHub Check: Build wheels for cp310-manylinux_aarch64
GitHub Check: Build wheels for cp311-macosx_x86_64
GitHub Check: Build wheels for cp311-macosx_arm64
GitHub Check: Build wheels for cp311-manylinux_x86_64
GitHub Check: Test Python (9, 3.10)
GitHub Check: Test C++ (true, true, true, false)
GitHub Check: Test Python (12, 3.13)
GitHub Check: Test Python (6, 3.13)
GitHub Check: Test Python (2, 3.10)
GitHub Check: Test C++ (false, true, true, false)
GitHub Check: Test C++ (true, false, false, true)
GitHub Check: Test Python (3, 3.13)
GitHub Check: Test C++ (false, false, false, true)
GitHub Check: Test Python (9, 3.13)
GitHub Check: Test Python (12, 3.10)
GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
GitHub Check: Test Python (11, 3.13)
GitHub Check: Test Python (11, 3.10)
GitHub Check: Test Python (3, 3.10)
GitHub Check: Test Python (5, 3.10)
GitHub Check: Test Python (8, 3.13)
GitHub Check: Test Python (7, 3.13)
GitHub Check: Test Python (4, 3.10)
GitHub Check: Test Python (10, 3.13)
GitHub Check: Test Python (10, 3.10)
GitHub Check: Test Python (7, 3.10)
GitHub Check: Test Python (6, 3.10)
GitHub Check: Test Python (2, 3.13)
GitHub Check: Test Python (4, 3.13)
GitHub Check: Test Python (1, 3.10)
GitHub Check: Test Python (5, 3.13)
GitHub Check: Test Python (1, 3.13)
GitHub Check: Test Python (8, 3.10)
GitHub Check: Analyze (python)
GitHub Check: Analyze (c-cpp)
GitHub Check: Build C++ (rocm, rocm)
GitHub Check: Build C++ (cpu, cpu)
GitHub Check: Build C++ (clang, clang)
GitHub Check: Build C++ (cuda120, cuda)

🔇 Additional comments (1)

deepmd/dpmodel/utils/region.py (1)

77-79: LGTM! Device and dtype consistency improved.

The explicit use of xp.ones((), dtype=icoord.dtype, device=array_api_compat.device(icoord)) ensures the remainder operation has matching dtype and device, preventing cross-device errors (especially in PyTorch) and maintaining numerical precision consistency across backends. This pattern is well-established throughout the codebase in env_mat_stat.py, nlist.py, and other utility modules.

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Copilot

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 EnvMatStat and EnvMatStatSe implementations with imports from dpmodel
Added build_type_exclude_mask method alias in PT's PairExcludeMask to 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.

…md-kit into pt-reuse-env-mat-stat

codecov · 2026-01-08T16:59:58Z

Codecov Report

❌ Patch coverage is 80.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.90%. Comparing base (8f021aa) to head (4c62a26).

Files with missing lines	Patch %	Lines
deepmd/dpmodel/utils/env_mat_stat.py	60.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5139      +/-   ##
==========================================
- Coverage   81.93%   81.90%   -0.03%     
==========================================
  Files         712      712              
  Lines       72864    72799      -65     
  Branches     3617     3616       -1     
==========================================
- Hits        59700    59628      -72     
- Misses      12001    12008       +7     
  Partials     1163     1163

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…iptor

njzjz added 3 commits January 8, 2026 19:23

refactor(pt): reuse env mat stat

0e5b449

sync from deepmodeling#5137

c895282

fix other errors

0521725

njzjz requested review from Copilot, iProzd and wanghan-iapcm January 8, 2026 11:37

github-actions bot added the Python label Jan 8, 2026

dosubot bot added the enhancement label Jan 8, 2026

[pre-commit.ci] auto fixes from pre-commit.com hooks

b652f03

for more information, see https://pre-commit.ci

Copilot started reviewing on behalf of njzjz January 8, 2026 11:40 View session

Copilot AI reviewed Jan 8, 2026

View reviewed changes

njzjz added 2 commits January 9, 2026 00:19

fix: improve remainder calculation in normalize_coord function

ad01332

Merge branch 'pt-reuse-env-mat-stat' of https://github.com/njzjz/deep…

c690b5d

…md-kit into pt-reuse-env-mat-stat

njzjz marked this pull request as draft January 8, 2026 18:18

njzjz added 4 commits January 10, 2026 14:42

Merge branch 'master' into pt-reuse-env-mat-stat

bc70c42

fix xp.ones

e72375e

fix IndexError: list index out of range

4058f9b

fix PairExcludeMask

a46cffd

njzjz removed request for iProzd and wanghan-iapcm January 10, 2026 07:57

njzjz and others added 2 commits January 10, 2026 17:01

add return

4c62a26

fix: validate last_dim in EnvMatStatSe for radial-only and full descr…

21b4c06

…iptor

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(pt): reuse dpmodel EnvMatStat #5139

refactor(pt): reuse dpmodel EnvMatStat #5139

njzjz commented Jan 8, 2026 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Jan 8, 2026 •

edited

Loading

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

Uh oh!

Copilot AI left a comment

Uh oh!

codecov bot commented Jan 8, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

refactor(pt): reuse dpmodel EnvMatStat #5139

Are you sure you want to change the base?

refactor(pt): reuse dpmodel EnvMatStat #5139

Conversation

njzjz commented Jan 8, 2026 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

codecov bot commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

njzjz commented Jan 8, 2026 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Jan 8, 2026 •

edited

Loading

codecov bot commented Jan 8, 2026 •

edited

Loading