BUG: cannot add take_along_axis to common/_aliases because of dask

Author: ev-br

array_api_compat/common/_aliases.py

@@ -517,12 +517,6 @@ def sort(
     return res
 
 
-# take_along_axis: axis defaults to -1; numpy, cupy do not have a default value;
-# pytorch defaults to None, which ravels.
-def take_along_axis(x: Array, indices: Array, /, *, xp: Namespace, axis: int = -1):
-    return xp.take_along_axis(x, indices, axis=axis)
-
-
 # nonzero should error for zero-dimensional arrays
 def nonzero(x: Array, /, xp: Namespace, **kwargs: object) -> tuple[Array, ...]:
     if x.ndim == 0:
@@ -719,7 +713,6 @@ def iinfo(type_: DType | Array, /, xp: Namespace) -> Any:
     "matmul",
     "matrix_transpose",
     "tensordot",
-    "take_along_axis",
     "vecdot",
     "isdtype",
     "unstack",

array_api_compat/cupy/_aliases.py

@@ -63,7 +63,6 @@
 sign = get_xp(cp)(_aliases.sign)
 finfo = get_xp(cp)(_aliases.finfo)
 iinfo = get_xp(cp)(_aliases.iinfo)
-take_along_axis = get_xp(cp)(_aliases.take_along_axis)
 
 _copy_default = object()
 
@@ -139,6 +138,11 @@ def count_nonzero(
    return result
 
 
+# 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 xp.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(cp, 'vecdot'):
@@ -160,6 +164,7 @@ def count_nonzero(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'count_nonzero', 'pow', 'sign']
+                              'bool', 'concat', 'count_nonzero', 'pow', 'sign',
+                              'take_along_axis']
 
 _all_ignore = ['cp', 'get_xp']

array_api_compat/numpy/_aliases.py

@@ -72,7 +72,6 @@
 sign = get_xp(np)(_aliases.sign)
 finfo = get_xp(np)(_aliases.finfo)
 iinfo = get_xp(np)(_aliases.iinfo)
-take_along_axis = get_xp(np)(_aliases.take_along_axis)
 
 
 def _supports_buffer_protocol(obj: object) -> TypeIs[Buffer]:  # pyright: ignore[reportUnusedFunction]
@@ -141,6 +140,11 @@ def count_nonzero(
     return result
 
 
+# take_along_axis: axis defaults to -1 but in numpy axis is a required arg
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
+    return xp.take_along_axis(x, indices, axis=axis)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(np, "vecdot"):
@@ -158,7 +162,6 @@ def count_nonzero(
 else:
     unstack = get_xp(np)(_aliases.unstack)
 
-
 __all__ = [
     "__array_namespace_info__",
     "asarray",
@@ -185,4 +188,3 @@ def count_nonzero(
 
 def __dir__() -> list[str]:
     return __all__
-