BUG: take_along_axis: numpy requires an axis

ev-br · ev-br · commit fdf9489054ee · 2025-04-21T14:59:13.000+02:00
diff --git a/array_api_compat/numpy/_aliases.py b/array_api_compat/numpy/_aliases.py
@@ -134,6 +134,13 @@ def count_nonzero(x: Array, axis=None, keepdims=False) -> Array:
     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'):
@@ -155,6 +162,7 @@ def count_nonzero(x: Array, axis=None, keepdims=False) -> Array:
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'count_nonzero', 'pow']
+                              'bool', 'concat', 'count_nonzero', 'pow',
+                              'take_along_axis']
 
 _all_ignore = ['np', 'get_xp']
