Add draft implementation of take_along_axis for 2024.12

As far as I can tell, NumPy matches the standard specification, except for the
fact that NumPy does not set a default value for axis.

Author: asmeurer

array_api_strict/__init__.py

@@ -262,9 +262,9 @@
     "trunc",
 ]
 
-from ._indexing_functions import take
+from ._indexing_functions import take, take_along_axis
 
-__all__ += ["take"]
+__all__ += ["take", "take_along_axis"]
 
 from ._info import __array_namespace_info__
 

array_api_strict/_indexing_functions.py

@@ -2,6 +2,7 @@
 
 from ._array_object import Array
 from ._dtypes import _integer_dtypes
+from ._flags import requires_api_version
 
 from typing import TYPE_CHECKING
 
@@ -25,3 +26,14 @@ def take(x: Array, indices: Array, /, *, axis: Optional[int] = None) -> Array:
     if x.device != indices.device:
         raise ValueError(f"Arrays from two different devices ({x.device} and {indices.device}) can not be combined.")
     return Array._new(np.take(x._array, indices._array, axis=axis), device=x.device)
+
+@requires_api_version('2024.12')
+def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.take_along_axis <numpy.take_along_axis>`.
+
+    See its docstring for more information.
+    """
+    if x.device != indices.device:
+        raise ValueError(f"Arrays from two different devices ({x.device} and {indices.device}) can not be combined.")
+    return Array._new(np.take_along_axis(x._array, indices._array, axis), device=x.device)

array_api_strict/tests/test_flags.py

@@ -284,12 +284,6 @@ def test_fft(func_name):
     'unstack': lambda: xp.unstack(xp.ones((3, 3)), axis=0),
 }
 
-api_version_2024_12_examples = {
-    'diff': lambda: xp.diff(xp.asarray([0, 1, 2])),
-    'nextafter': lambda: xp.nextafter(xp.asarray(0.), xp.asarray(1.)),
-    'reciprocal': lambda: xp.reciprocal(xp.asarray([2.])),
-}
-
 @pytest.mark.parametrize('func_name', api_version_2023_12_examples.keys())
 def test_api_version_2023_12(func_name):
     func = api_version_2023_12_examples[func_name]
@@ -308,6 +302,14 @@ def test_api_version_2023_12(func_name):
     set_array_api_strict_flags(api_version='2022.12')
     pytest.raises(RuntimeError, func)
 
+api_version_2024_12_examples = {
+    'diff': lambda: xp.diff(xp.asarray([0, 1, 2])),
+    'nextafter': lambda: xp.nextafter(xp.asarray(0.), xp.asarray(1.)),
+    'reciprocal': lambda: xp.reciprocal(xp.asarray([2.])),
+    'take_along_axis': lambda: xp.take_along_axis(xp.zeros((2, 3)),
+                                                  xp.zeros((1, 4), dtype=xp.int64)),
+}
+
 @pytest.mark.parametrize('func_name', api_version_2024_12_examples.keys())
 def test_api_version_2024_12(func_name):
     func = api_version_2024_12_examples[func_name]