style: extend device lint check to Array API

[njzjz]

Summary by CodeRabbit

• Bug Fixes
  • Fixed device placement in tensor operations so tensors are allocated on the same device as inputs, avoiding cross-device errors.
  • Improved device inference for matrix/identity operations to ensure computations occur on the correct device.
  • Expanded device validation to consistently cover all supported compute backends and surface missing-device issues.

[Copilot]

Pull request overview

Extends the repository’s “explicit device” linting and updates Array API-based tensor creation sites to consistently place newly created arrays on the intended device (avoiding accidental CPU allocations when running with torch-backed Array API).

Changes:

• Extend no-explicit-device lint rule to cover xp.* creation calls in addition to torch.*.
• Propagate device explicitly when creating xp.eye(...) in the polar atomic model output-stat path.
• Add device=... to Array API helper allocations (xp.arange, xp.zeros) to match the new lint requirement.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
source/checker/deepmd_checker.py	Lint rule now flags missing device= for xp.* creation ops (in addition to torch.*).
deepmd/dpmodel/atomic_model/polar_atomic_model.py	Ensures xp.eye is created on the same device as the model’s bias tensors.
deepmd/dpmodel/array_api.py	Adds explicit device= when creating intermediate arrays in array-api helper utilities.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c468ed6 and 20b055b.

📒 Files selected for processing (1)

• deepmd/dpmodel/array_api.py

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/dpmodel/array_api.py

---

📝 Walkthrough

Walkthrough

The PR updates device placement specifications across array API operations and adds device inference logic in model code. It also expands device-check validation in the linter to include both PyTorch and XP backends, ensuring tensors are created on consistent devices.

Changes

Cohort / File(s)	Summary
Device Placement in Array API deepmd/dpmodel/array_api.py	Updated xp_scatter_sum and xp_bincount to pass device=array_api_compat.device(...) when creating xp.arange/xp.zeros, ensuring tensors are allocated on the input/weights device; formatting adjusted.
Device Inference in Model deepmd/dpmodel/atomic_model/polar_atomic_model.py	Added device inference from out_bias[self.bias_keys[0]] and used device=device when constructing the 3x3 identity matrix used for bias/diagonal adjustments.
Backend Validation Expansion source/checker/deepmd_checker.py	Broadened device-check in visit_call to treat both "torch" and "xp" backends the same for missing-device diagnostics by checking node.func.expr.name in {"torch", "xp"}.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• style: extend no-explicit-dtype check to xp and jnp #4247: Updates the same linter/device-check logic to include "xp" alongside "torch" in backend checks for visit_call.

Suggested reviewers

• wanghan-iapcm

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'style: extend device lint check to Array API' directly reflects the main change: extending the device lint check (in deepmd_checker.py) to cover both torch and xp backends.

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

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@deepmd/dpmodel/array_api.py`:
- Around line 81-83: The call to array_api_compat.device(...) is passing
input.device which will raise AttributeError for NumPy arrays; change the
argument to pass the array itself (input) instead. Locate the xp.arange call
that constructs idx (the line using idx = xp.arange(...
device=array_api_compat.device(...))) and replace
array_api_compat.device(input.device) with array_api_compat.device(input) so it
matches the other usages across the codebase and works for both NumPy and
device-backed arrays.

In `@source/checker/deepmd_checker.py`:
- Around line 63-65: Update the misleading comment and lint text that claim
"only PT needs device": since the check now triggers when node.func.expr.name is
in {"torch", "xp"}, change the inline comment near the no_device check to
mention both PyTorch and the xp backend, and update the lint message/description
associated with the "no-explicit-device" rule (or its registration) so its
wording clearly covers both torch and xp backends rather than implying
PyTorch-only; look for uses of no_device, node.func.expr.name and the
self.add_message("no-explicit-device", node=node) call and adjust the comment
and rule message/description accordingly.

---

ℹ️ Review info

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 65eea4b and c468ed6.

📒 Files selected for processing (3)

• deepmd/dpmodel/array_api.py
• deepmd/dpmodel/atomic_model/polar_atomic_model.py
• source/checker/deepmd_checker.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.00%. Comparing base (65eea4b) to head (20b055b).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5261   +/-   ##
=======================================
  Coverage   82.00%   82.00%           
=======================================
  Files         750      750           
  Lines       75082    75083    +1     
  Branches     3615     3616    +1     
=======================================
+ Hits        61571    61572    +1     
  Misses      12347    12347           
  Partials     1164     1164           

☔ 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.