MAINT: move take_along_axis to common/_aliases

Author: ev-br

array_api_compat/common/_aliases.py

@@ -517,6 +517,12 @@ 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:
@@ -713,6 +719,7 @@ 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,6 +63,7 @@
 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()
 

array_api_compat/numpy/_aliases.py

@@ -72,6 +72,7 @@
 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]
@@ -140,13 +141,6 @@ def count_nonzero(
     return result
 
 
-# "axis=-1" is an optional argument of `take_along_axis` but numpy has no default
-def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
-    if axis is None:
-        axis = -1
-    return np.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"):