mv methods to array api

Author: iProzd

deepmd/dpmodel/array_api.py

@@ -92,3 +92,37 @@ def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray
         )
     else:
         raise NotImplementedError("Only JAX arrays are supported.")
+
+
+def add_at(x, indices, values):
+    """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)
+    if array_api_compat.is_numpy_array(x):
+        # NumPy: supports np.add.at (in-place)
+        xp.add.at(x, indices, values)
+        return x
+
+    elif array_api_compat.is_jax_array(x):
+        # JAX: functional update, not in-place
+        return x.at[indices].add(values)
+    else:
+        # Fallback for array_api_strict: use basic indexing only
+        # may need a more efficient way to do this
+        n = indices.shape[0]
+        for i in range(n):
+            idx = indices[i]
+            x[idx] = x[idx] + values[i]
+        return x
+
+
+def bincount(x, weights=None, minlength=0):
+    """Counts the number of occurrences of each value in x."""
+    xp = array_api_compat.array_namespace(x)
+    if array_api_compat.is_numpy_array(x) or array_api_compat.is_jax_array(x):
+        result = xp.bincount(x, weights=weights, minlength=minlength)
+    else:
+        if weights is None:
+            weights = xp.ones_like(x)
+        result = xp.zeros((max(minlength, int(xp.max(x)) + 1),), dtype=weights.dtype)
+        result = add_at(result, x, weights)
+    return result

deepmd/dpmodel/utils/network.py

@@ -21,6 +21,8 @@
     NativeOP,
 )
 from deepmd.dpmodel.array_api import (
+    add_at,
+    bincount,
     support_array_api,
 )
 from deepmd.dpmodel.common import (
@@ -983,42 +985,15 @@ def aggregate(
     output: [num_owner, feature_dim]
     """
     xp = array_api_compat.array_namespace(data, owners)
-
-    def add_at(x, indices, values):
-        unique_ids = xp.unique(indices)
-        for i in unique_ids:
-            mask = xp.where(indices == i, 1, 0)
-            mask = xp.expand_dims(mask, axis=1) if len(values.shape) != 1 else mask
-            selected = values * mask
-            summed = xp.sum(selected, axis=0)
-            x[i] = x[i] + summed
-
-        return x
-
-    def bincount(x, weights=None, minlength=0):
-        if weights is None:
-            weights = xp.ones_like(x)
-        result = xp.zeros((max(minlength, int(xp.max(x)) + 1),), dtype=weights.dtype)
-        result = add_at(result, x, weights)
-        return result
-
-    if hasattr(xp, "bincount"):
-        bin_count = xp.bincount(owners)
-    else:
-        # for array_api_strict
-        bin_count = bincount(owners)
+    bin_count = bincount(owners)
     bin_count = xp.where(bin_count == 0, xp.ones_like(bin_count), bin_count)
 
     if num_owner is not None and bin_count.shape[0] != num_owner:
         difference = num_owner - bin_count.shape[0]
         bin_count = xp.concat([bin_count, xp.ones(difference, dtype=bin_count.dtype)])
 
     output = xp.zeros((bin_count.shape[0], data.shape[1]), dtype=data.dtype)
-    if hasattr(xp, "add") and hasattr(xp.add, "at"):
-        xp.add.at(output, owners, data)
-    else:
-        # for array_api_strict
-        output = add_at(output, owners, data)
+    output = add_at(output, owners, data)
 
     if average:
         output = xp.transpose(xp.transpose(output) / bin_count)