feat(dpmodel): add PyTorch support to array_api utilities

[Copilot]

The array_api.py module only supported JAX and NumPy arrays. PyTorch tensors now work with all array API utilities.

Changes

• xp_scatter_sum: Added PyTorch path using torch.scatter_add()
• xp_add_at: Added PyTorch path using torch.index_add()
• xp_bincount: Extended backend check to include PyTorch arrays

All PyTorch operations use non-mutating variants (no trailing _) to maintain functional semantics consistent with JAX.

Testing

• Created comprehensive consistency tests in source/tests/consistent/test_array_api.py
• Tests verify consistency across NumPy, PyTorch, JAX, and array_api_strict backends
• Tests verify non-mutating behavior for PyTorch and JAX
• All 10 tests pass successfully (5 JAX tests skipped when JAX not installed)

Usage

import torch
from deepmd.dpmodel.array_api import xp_bincount, xp_add_at

# PyTorch tensors now work seamlessly
x = torch.tensor([0, 1, 1, 3, 2, 1, 7])
counts = xp_bincount(x)  # tensor([1, 3, 1, 1, 0, 0, 0, 1])

# Operations are non-mutating
x = torch.zeros(5, 3)
result = xp_add_at(x, torch.tensor([0, 1]), torch.ones(2, 3))
# x remains unchanged, result contains the update

Original prompt

The current deepmd/dpmodel/array_api.py only considers JAX and NumPy. Please also support pytorch.

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

[codecov]

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 81.96%. Comparing base (6112545) to head (9c45114).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/dpmodel/array_api.py	87.50%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5198   +/-   ##
=======================================
  Coverage   81.96%   81.96%           
=======================================
  Files         714      714           
  Lines       73502    73508    +6     
  Branches     3615     3616    +1     
=======================================
+ Hits        60247    60252    +5     
  Misses      12092    12092           
- Partials     1163     1164    +1     

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

[Copilot]

Pull request overview

Adds PyTorch support to deepmd/dpmodel/array_api.py utility functions so DPModel’s array-API helpers can operate on torch tensors in addition to existing NumPy/JAX support, with new consistency tests to validate behavior across backends.

Changes:

• Added PyTorch implementations for xp_scatter_sum (torch.scatter_add) and xp_add_at (torch.index_add).
• Extended xp_bincount backend handling to include PyTorch arrays.
• Introduced new consistency tests covering PyTorch/JAX/NumPy (and array_api_strict where applicable).

Reviewed changes

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

File	Description
deepmd/dpmodel/array_api.py	Adds torch branches for scatter/add-at utilities and enables torch in xp_bincount.
source/tests/consistent/test_array_api.py	New backend-consistency tests for xp_scatter_sum, xp_add_at, and xp_bincount including non-mutating checks for torch.

Comments suppressed due to low confidence (1)

deepmd/dpmodel/array_api.py:113

• xp_add_at docstring mentions returning a new array only for JAX, but the new PyTorch implementation is also out-of-place (non-mutating). Update the docstring to reflect that behavior for PyTorch as well, so callers don’t assume PyTorch will mutate like NumPy.

def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
    """Adds values to the specified indices of x in place or returns new x (for JAX)."""
    xp = array_api_compat.array_namespace(x, indices, values)

---

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