refactor: consolidate backend logic into array_api.py with generic implementations (#5202)

## Refactor backend logic into array_api.py

### Analysis Complete
- [x] Identify all JAX-specific code in deepmd/dpmodel/
- [x] Review existing array_api.py structure
- [x] Identify test infrastructure

### Implementation Complete
- [x] Add `xp_sigmoid` function to array_api.py for backend-specific
sigmoid handling
- [x] Add `xp_setitem_at` function to array_api.py for backend-specific
array assignment
- [x] Implement generic `xp_scatter_sum` using array_api operations
- [x] Update network.py to use new `xp_sigmoid` function
- [x] Update network.py to use new `xp_setitem_at` function
- [x] Update transform_output.py to use `xp_scatter_sum` from
array_api.py
- [x] Add comprehensive tests for all functions
- [x] Add array-api-strict tests for complete backend coverage
- [x] Remove unused JAX-specific scatter_sum implementation
- [x] Make xp_setitem_at non-mutating for PyTorch (consistent with other
xp_* helpers)
- [x] Run tests to validate changes (20 passed, 7 skipped)
- [x] Run linters to ensure code quality (all checks passed)

### Changes Summary

1. **deepmd/dpmodel/array_api.py** - Centralized backend
implementations:
- **xp_sigmoid**: Added PyTorch-specific `torch.sigmoid()` and
JAX-specific `jax.nn.sigmoid()` implementations with generic fallback
- **xp_setitem_at**: Handles JAX's functional `.at[].set()` syntax and
PyTorch's `clone()` for non-mutating behavior, while NumPy uses in-place
assignment
- **xp_scatter_sum**: Implemented generic array_api version using
`xp_take_along_axis` and `xp_add_at` helper functions (merged JAX
implementation logic), with PyTorch optimization retained
   - Now supports NumPy, JAX, PyTorch, and array-api-strict backends
- **All xp_* helpers are now consistently non-mutating for PyTorch and
JAX**

2. **deepmd/dpmodel/utils/network.py** - Refactored to use centralized
functions:
- Replaced direct JAX conditional in `sigmoid_t` function with call to
`xp_sigmoid`
- Replaced JAX conditional for array assignment with call to
`xp_setitem_at`

3. **deepmd/dpmodel/model/transform_output.py** - Simplified scatter_sum
usage:
- Removed JAX conditional and direct import of
`deepmd.jax.common.scatter_sum`
   - Now uses `xp_scatter_sum` from array_api.py consistently

4. **deepmd/jax/common.py** - Removed unused code:
- Removed `scatter_sum` function (replaced by generic implementation in
array_api.py)
   - Function was not imported or used anywhere in the codebase

5. **source/tests/consistent/test_array_api.py** - Complete test
coverage across all backends:
   - Added array-api-strict tests for `TestXpSigmoidConsistent`
   - Added array-api-strict tests for `TestXpSetitemAtConsistent`
   - Added array-api-strict tests for `TestXpScatterSumConsistent`
   - Added array-api-strict tests for `TestXpBincountConsistent`
- Added array-api-strict tests for `TestXpBincountWithWeightsConsistent`
- Added array-api-strict tests for
`TestXpBincountWithMinlengthConsistent`
- **Added non-mutating verification for PyTorch in
`TestXpSetitemAtConsistent`** (consistent with other tests)
- All test classes now have consistent coverage across NumPy, PyTorch,
JAX, and array-api-strict

### Test Results
- All 27 array_api tests pass (20 passed, 7 skipped for JAX not
installed)
- Complete array-api-strict test coverage for all xp_* functions
- All PyTorch tests verify non-mutating behavior
- All network tests pass
- Ruff linting: all checks passed

### Benefits
- ✅ All backend-specific conditionals (`if is_jax_array`, `if
is_torch_array`) centralized in array_api.py
- ✅ No direct backend imports (`from deepmd.jax`, `from torch`) outside
of array_api.py
- ✅ Generic array_api implementations support all backends (NumPy, JAX,
PyTorch, array-api-strict)
- ✅ Backend-specific optimizations retained where beneficial (e.g.,
`torch.sigmoid`, `torch.scatter_add`)
- ✅ **All xp_* helpers are consistently non-mutating for PyTorch and
JAX**, preventing autograd issues
- ✅ Removed duplicate/unused code from JAX backend
- ✅ Improved maintainability - future backend changes only need to
update array_api.py
- ✅ Consistent API for backend-specific operations across the codebase

<!-- START COPILOT ORIGINAL PROMPT -->



<details>

<summary>Original prompt</summary>

> 
> ----
> 
> *This section details on the original issue you should resolve*
> 
> <issue_title>[Feature Request] Refactor JAX backend logic in
deepmd/dpmodel into array_api.py</issue_title>
> <issue_description>### Summary
> 
> Refactor all JAX-specific backend code currently spread across files
in deepmd/dpmodel/ (except for array_api.py) by consolidating these
implementations into deepmd/dpmodel/array_api.py. This will reduce code
duplication, improve maintainability, and centralize NUMPY, JAX, and
TORCH backend logic in one place.
> 
> ### Detailed Description
> 
> Many files in deepmd/dpmodel/ contain implementation blocks like `if
is_jax_array` for JAX backend support. To streamline future backend
improvements and reduce maintenance overhead, all such JAX-specific
logic should be migrated into array_api.py. This centralization will
make backend wrappers easier to maintain and update.
> 
> Example JAX references found:
> - deepmd/dpmodel/utils/network.py, ~line 1149: JAX-specific array
assignment logic
> - deepmd/dpmodel/model/transform_output.py, ~line 217: JAX scatter_sum
usage
> 
> After the refactoring, other files should invoke unified backend
functions from array_api.py instead of handling JAX conditionals
locally.
> 
> (Search results are limited, please see [GitHub
search](https://github.com/deepmodeling/deepmd-kit/search?q=is_jax_array+path%3Adeepmd%2Fdpmodel%2F)
for full context.)
> 
> ### Further Information, Files, and Links
> 
> _No response_</issue_description>
> 
> ## Comments on the Issue (you are @copilot in this section)
> 
> <comments>
> </comments>
> 


</details>



<!-- START COPILOT CODING AGENT SUFFIX -->

- Fixes #5201

<!-- START COPILOT CODING AGENT TIPS -->
---

💡 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](https://gh.io/copilot-coding-agent-tips) in the docs.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/dpmodel/array_api.py

@@ -57,25 +57,38 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
 
 
 def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array:
-    """Reduces all values from the src tensor to the indices specified in the index tensor."""
-    if array_api_compat.is_jax_array(input):
-        from deepmd.jax.common import (
-            scatter_sum,
-        )
+    """Reduces all values from the src tensor to the indices specified in the index tensor.
 
-        return scatter_sum(
-            input,
-            dim,
-            index,
-            src,
-        )
-    elif array_api_compat.is_torch_array(input):
-        # PyTorch: use scatter_add (non-mutating version)
+    This function is similar to PyTorch's scatter_add and JAX's scatter_sum.
+    It adds values from src to input at positions specified by index along the given dimension.
+    """
+    if array_api_compat.is_torch_array(input):
+        # PyTorch: use scatter_add (non-mutating version) for better performance
         import torch
 
         return torch.scatter_add(input, dim, index, src)
-    else:
-        raise NotImplementedError("Only JAX and PyTorch arrays are supported.")
+
+    # Generic array_api implementation (works for JAX, NumPy, array-api-strict, etc.)
+    xp = array_api_compat.array_namespace(input)
+
+    # Create flat index array matching input shape
+    idx = xp.arange(input.size, dtype=xp.int64)
+    idx = xp.reshape(idx, input.shape)
+
+    # Get flat indices where we want to add values
+    new_idx = xp_take_along_axis(idx, index, axis=dim)
+    new_idx = xp.reshape(new_idx, (-1,))
+
+    # Flatten arrays
+    shape = input.shape
+    input_flat = xp.reshape(input, (-1,))
+    src_flat = xp.reshape(src, (-1,))
+
+    # Add values at the specified indices
+    result = xp_add_at(input_flat, new_idx, src_flat)
+
+    # Reshape back to original shape
+    return xp.reshape(result, shape)
 
 
 def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
@@ -104,6 +117,61 @@ def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
         return x
 
 
+def xp_sigmoid(x: Array) -> Array:
+    """Compute the sigmoid function.
+
+    JAX and PyTorch have optimized sigmoid implementations.
+    See https://github.com/jax-ml/jax/discussions/15617
+    """
+    if array_api_compat.is_jax_array(x):
+        from deepmd.jax.env import (
+            jax,
+        )
+
+        return jax.nn.sigmoid(x)
+    elif array_api_compat.is_torch_array(x):
+        import torch
+
+        return torch.sigmoid(x)
+    xp = array_api_compat.array_namespace(x)
+    return 1 / (1 + xp.exp(-x))
+
+
+def xp_setitem_at(x: Array, mask: Array, values: Array) -> Array:
+    """Set items at boolean mask indices.
+
+    For JAX and PyTorch arrays, returns a new array (non-mutating).
+    For NumPy arrays, modifies in-place and returns the same array.
+
+    Parameters
+    ----------
+    x : Array
+        The array to modify
+    mask : Array
+        Boolean mask indicating positions to set
+    values : Array
+        Values to set at masked positions
+
+    Returns
+    -------
+    Array
+        Modified array (new array for JAX/PyTorch, same array for NumPy)
+    """
+    if array_api_compat.is_jax_array(x):
+        # JAX doesn't support in-place item assignment
+        return x.at[mask].set(values)
+    elif array_api_compat.is_torch_array(x):
+        # PyTorch: clone to avoid mutating the input (non-mutating version)
+        import torch
+
+        result = torch.clone(x)
+        result[mask] = values
+        return result
+    # Standard item assignment for NumPy, array-api-strict, etc.
+    x[mask] = values
+    return x
+
+
 def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> Array:
     """Counts the number of occurrences of each value in x."""
     xp = array_api_compat.array_namespace(x)

deepmd/dpmodel/model/transform_output.py

@@ -219,20 +219,12 @@ def communicate_extended_output(
                         vldims + derv_c_ext_dims,
                         dtype=vv.dtype,
                     )
-                    # jax only
-                    if array_api_compat.is_jax_array(virial):
-                        from deepmd.jax.common import (
-                            scatter_sum,
-                        )
-
-                        virial = scatter_sum(
-                            virial,
-                            1,
-                            mapping,
-                            model_ret[kk_derv_c],
-                        )
-                    else:
-                        raise NotImplementedError("Only JAX arrays are supported.")
+                    virial = xp_scatter_sum(
+                        virial,
+                        1,
+                        mapping,
+                        model_ret[kk_derv_c],
+                    )
                     new_ret[kk_derv_c] = virial
                     new_ret[kk_derv_c + "_redu"] = xp.sum(new_ret[kk_derv_c], axis=1)
                 else:

deepmd/dpmodel/utils/network.py

@@ -25,6 +25,8 @@
     Array,
     xp_add_at,
     xp_bincount,
+    xp_setitem_at,
+    xp_sigmoid,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -39,15 +41,7 @@
 
 def sigmoid_t(x):  # noqa: ANN001, ANN201
     """Sigmoid."""
-    if array_api_compat.is_jax_array(x):
-        from deepmd.jax.env import (
-            jax,
-        )
-
-        # see https://github.com/jax-ml/jax/discussions/15617
-        return jax.nn.sigmoid(x)
-    xp = array_api_compat.array_namespace(x)
-    return 1 / (1 + xp.exp(-x))
+    return xp_sigmoid(x)
 
 
 class Identity(NativeOP):
@@ -1251,11 +1245,7 @@ def get_graph_index(  # noqa: ANN201
     # edge(ij) to angle(ijk) index_select; angle(ijk) to edge(ij) aggregate
     edge_id = xp.arange(n_edge, dtype=nlist.dtype)
     edge_index = xp.zeros((nf, nloc, nnei), dtype=nlist.dtype)
-    if array_api_compat.is_jax_array(nlist):
-        # JAX doesn't support in-place item assignment
-        edge_index = edge_index.at[xp.astype(nlist_mask, xp.bool)].set(edge_id)
-    else:
-        edge_index[xp.astype(nlist_mask, xp.bool)] = edge_id
+    edge_index = xp_setitem_at(edge_index, xp.astype(nlist_mask, xp.bool), edge_id)
     # only cut a_nnei neighbors, to avoid nnei x nnei
     edge_index = edge_index[:, :, :a_nnei]
     edge_index_ij = xp.broadcast_to(

deepmd/jax/common.py

@@ -94,15 +94,3 @@ def __dlpack__(self, *args: Any, **kwargs: Any) -> Any:
 
     def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any:
         return self.value.__dlpack_device__(*args, **kwargs)
-
-
-def scatter_sum(
-    input: jnp.ndarray, dim: int, index: jnp.ndarray, src: jnp.ndarray
-) -> jnp.ndarray:
-    """Reduces all values from the src tensor to the indices specified in the index tensor."""
-    idx = jnp.arange(input.size, dtype=jnp.int64).reshape(input.shape)
-    new_idx = jnp.take_along_axis(idx, index, axis=dim).ravel()
-    shape = input.shape
-    input = input.ravel()
-    input = input.at[new_idx].add(src.ravel())
-    return input.reshape(shape)

source/tests/consistent/test_array_api.py

@@ -8,6 +8,8 @@
     xp_add_at,
     xp_bincount,
     xp_scatter_sum,
+    xp_setitem_at,
+    xp_sigmoid,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -66,6 +68,19 @@ def test_jax_consistent_with_ref(self) -> None:
         result = xp_scatter_sum(input_jax, self.dim, index_jax, src_jax)
         np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
 
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        input_xp = xp.asarray(self.input_np)
+        index_xp = xp.asarray(self.index_np)
+        src_xp = xp.asarray(self.src_np)
+        result = xp_scatter_sum(input_xp, self.dim, index_xp, src_xp)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
 
 class TestXpAddAtConsistent(unittest.TestCase):
     """Test xp_add_at consistency across backends."""
@@ -139,6 +154,17 @@ def test_jax_consistent_with_ref(self) -> None:
         result = xp_bincount(x_jax)
         np.testing.assert_equal(self.ref, to_numpy_array(result))
 
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        x_xp = xp.asarray(self.x_np)
+        result = xp_bincount(x_xp)
+        np.testing.assert_equal(self.ref, to_numpy_array(result))
+
 
 class TestXpBincountWithWeightsConsistent(unittest.TestCase):
     """Test xp_bincount with weights consistency across backends."""
@@ -166,6 +192,18 @@ def test_jax_consistent_with_ref(self) -> None:
         result = xp_bincount(x_jax, weights=weights_jax)
         np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
 
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        x_xp = xp.asarray(self.x_np)
+        weights_xp = xp.asarray(self.weights_np)
+        result = xp_bincount(x_xp, weights=weights_xp)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
 
 class TestXpBincountWithMinlengthConsistent(unittest.TestCase):
     """Test xp_bincount with minlength consistency across backends."""
@@ -190,3 +228,99 @@ def test_jax_consistent_with_ref(self) -> None:
         x_jax = jnp.array(self.x_np)
         result = xp_bincount(x_jax, minlength=self.minlength)
         np.testing.assert_equal(self.ref, to_numpy_array(result))
+
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        x_xp = xp.asarray(self.x_np)
+        result = xp_bincount(x_xp, minlength=self.minlength)
+        np.testing.assert_equal(self.ref, to_numpy_array(result))
+
+
+class TestXpSigmoidConsistent(unittest.TestCase):
+    """Test xp_sigmoid consistency across backends."""
+
+    def setUp(self) -> None:
+        self.x_np = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
+        # Reference using NumPy sigmoid
+        self.ref = 1 / (1 + np.exp(-self.x_np))
+
+    def test_numpy_consistent_with_ref(self) -> None:
+        result = xp_sigmoid(self.x_np)
+        np.testing.assert_allclose(self.ref, result, atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_pt_consistent_with_ref(self) -> None:
+        x_pt = torch.from_numpy(self.x_np)
+        result = xp_sigmoid(x_pt)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_JAX, "JAX is not installed")
+    def test_jax_consistent_with_ref(self) -> None:
+        x_jax = jnp.array(self.x_np)
+        result = xp_sigmoid(x_jax)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        x_xp = xp.asarray(self.x_np)
+        result = xp_sigmoid(x_xp)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
+
+class TestXpSetitemAtConsistent(unittest.TestCase):
+    """Test xp_setitem_at consistency across backends."""
+
+    def setUp(self) -> None:
+        self.x_np = np.zeros((5, 3))
+        self.mask_np = np.array([True, False, True, False, True])
+        self.values_np = np.ones((3, 3))
+        # Reference using NumPy
+        self.ref = self.x_np.copy()
+        self.ref[self.mask_np] = self.values_np
+
+    def test_numpy_consistent_with_ref(self) -> None:
+        x = self.x_np.copy()
+        result = xp_setitem_at(x, self.mask_np, self.values_np)
+        np.testing.assert_allclose(self.ref, result, atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_pt_consistent_with_ref(self) -> None:
+        x_pt = torch.from_numpy(self.x_np)
+        mask_pt = torch.from_numpy(self.mask_np)
+        values_pt = torch.from_numpy(self.values_np)
+        result = xp_setitem_at(x_pt, mask_pt, values_pt)
+        # Verify original tensor is unchanged (non-mutating)
+        np.testing.assert_allclose(self.x_np, to_numpy_array(x_pt), atol=1e-10)
+        # Verify result matches reference
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
+    @unittest.skipUnless(INSTALLED_JAX, "JAX is not installed")
+    def test_jax_consistent_with_ref(self) -> None:
+        x_jax = jnp.array(self.x_np)
+        mask_jax = jnp.array(self.mask_np)
+        values_jax = jnp.array(self.values_np)
+        result = xp_setitem_at(x_jax, mask_jax, values_jax)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)
+
+    @unittest.skipUnless(
+        INSTALLED_ARRAY_API_STRICT, "array_api_strict is not installed"
+    )
+    @unittest.skipUnless(
+        sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
+    )
+    def test_array_api_strict_consistent_with_ref(self) -> None:
+        x_xp = xp.asarray(self.x_np)
+        mask_xp = xp.asarray(self.mask_np)
+        values_xp = xp.asarray(self.values_np)
+        result = xp_setitem_at(x_xp, mask_xp, values_xp)
+        np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10)