rename

Author: iProzd

deepmd/dpmodel/array_api.py

@@ -94,7 +94,7 @@ def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray
         raise NotImplementedError("Only JAX arrays are supported.")
 
 
-def add_at(x, indices, values):
+def xp_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):
@@ -115,7 +115,7 @@ def add_at(x, indices, values):
         return x
 
 
-def bincount(x, weights=None, minlength=0):
+def xp_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):
@@ -124,5 +124,5 @@ 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)
+        result = xp_add_at(result, x, weights)
     return result

deepmd/dpmodel/utils/network.py

@@ -21,9 +21,9 @@
     NativeOP,
 )
 from deepmd.dpmodel.array_api import (
-    add_at,
-    bincount,
     support_array_api,
+    xp_add_at,
+    xp_bincount,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -985,15 +985,15 @@ def aggregate(
     output: [num_owner, feature_dim]
     """
     xp = array_api_compat.array_namespace(data, owners)
-    bin_count = bincount(owners)
+    bin_count = xp_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)
-    output = add_at(output, owners, data)
+    output = xp_add_at(output, owners, data)
 
     if average:
         output = xp.transpose(xp.transpose(output) / bin_count)