ENH: test take_along_axis

Author: ev-br

array_api_tests/test_indexing_functions.py

@@ -60,3 +60,55 @@ def test_take(x, data):
     # sanity check
     with pytest.raises(StopIteration):
         next(out_indices)
+
+
+
+@pytest.mark.unvectorized
+@pytest.mark.min_version("2024.12")
+@given(
+    x=hh.arrays(hh.all_dtypes, hh.shapes(min_dims=1, min_side=1)),
+    data=st.data(),
+)
+def test_take_along_axis(x, data):
+    # TODO
+    # 1. negative axis
+    # 2. negative indices
+    # 3. different dtypes for indices
+    axis = data.draw(st.integers(0, max(x.ndim - 1, 0)), label="axis")
+    len_axis = data.draw(st.integers(0, 2*x.shape[axis]), label="len_axis")
+
+    idx_shape = x.shape[:axis] + (len_axis,) + x.shape[axis+1:]
+    indices = data.draw(
+        hh.arrays(
+            shape=idx_shape,
+            dtype=dh.default_int,
+            elements={"min_value": 0, "max_value": x.shape[axis]-1}
+        ),
+        label="indices"
+    )
+    note(f"{indices=}  {idx_shape=}")
+
+    out = xp.take_along_axis(x, indices, axis=axis)
+
+    ph.assert_dtype("take_along_axis", in_dtype=x.dtype, out_dtype=out.dtype)
+    ph.assert_shape(
+        "take_along_axis",
+        out_shape=out.shape,
+        expected=x.shape[:axis] + (len_axis,) + x.shape[axis+1:],
+        kw=dict(
+            x=x,
+            indices=indices,
+            axis=axis,
+        ),
+    )
+
+    # value test: notation is from `np.take_along_axis` docstring
+    Ni, Nk = x.shape[:axis], x.shape[axis+1:]
+    for ii in sh.ndindex(Ni):
+        for kk in sh.ndindex(Nk):
+            a_1d = x[ii + (slice(None),) + kk]
+            i_1d = indices[ii + (slice(None),) + kk]
+            o_1d = out[ii + (slice(None),) + kk]
+            for j in range(len_axis):
+                assert o_1d[j] == a_1d[i_1d[j]], f'{ii=}, {kk=}, {j=}'
+