..

Author: ev-br

array_api_compat/cupy/_aliases.py

@@ -140,7 +140,7 @@ def count_nonzero(
 
 # take_along_axis: axis defaults to -1 but in cupy (and numpy) axis is a required arg
 def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
-    return np.take_along_axis(x, indices, axis=axis)
+    return cp.take_along_axis(x, indices, axis=axis)
 
 
 # These functions are completely new here. If the library already has them