test_take

Author: honno

array_api_tests/_array_module.py

@@ -62,6 +62,7 @@ def __repr__(self):
 ]
 _constants = ["e", "inf", "nan", "pi"]
 _funcs = [f.__name__ for funcs in stubs.category_to_funcs.values() for f in funcs]
+_funcs += ["take"]  # TODO: bump spec and update array-api-tests to new spec layout
 _top_level_attrs = _dtypes + _constants + _funcs + stubs.EXTENSIONS
 
 for attr in _top_level_attrs:

array_api_tests/test_indexing_functions.py

@@ -0,0 +1,62 @@
+import pytest
+from hypothesis import given, note
+from hypothesis import strategies as st
+
+from . import _array_module as xp
+from . import dtype_helpers as dh
+from . import hypothesis_helpers as hh
+from . import pytest_helpers as ph
+from . import shape_helpers as sh
+from . import xps
+
+pytestmark = pytest.mark.ci
+
+
+@pytest.mark.min_version("2022.12")
+@given(
+    x=xps.arrays(xps.scalar_dtypes(), hh.shapes(min_dims=1, min_side=1)),
+    data=st.data(),
+)
+def test_take(x, data):
+    # TODO:
+    # * negative axis
+    # * negative indices
+    # * different dtypes for indices
+    axis = data.draw(st.integers(0, max(x.ndim - 1, 0)), label="axis")
+    _indices = data.draw(
+        st.lists(st.integers(0, x.shape[axis] - 1), min_size=1, unique=True),
+        label="_indices",
+    )
+    indices = xp.asarray(_indices, dtype=dh.default_int)
+    note(f"{indices=}")
+
+    out = xp.take(x, indices, axis=axis)
+
+    ph.assert_dtype("take", x.dtype, out.dtype)
+    ph.assert_shape(
+        "take",
+        out.shape,
+        x.shape[:axis] + (len(_indices),) + x.shape[axis + 1 :],
+        x=x,
+        indices=indices,
+        axis=axis,
+    )
+    out_indices = sh.ndindex(out.shape)
+    axis_indices = list(sh.axis_ndindex(x.shape, axis))
+    for axis_idx in axis_indices:
+        f_axis_idx = sh.fmt_idx("x", axis_idx)
+        for i in _indices:
+            f_take_idx = sh.fmt_idx(f_axis_idx, i)
+            indexed_x = x[axis_idx][i]
+            for at_idx in sh.ndindex(indexed_x.shape):
+                out_idx = next(out_indices)
+                ph.assert_0d_equals(
+                    "take",
+                    sh.fmt_idx(f_take_idx, at_idx),
+                    indexed_x[at_idx],
+                    sh.fmt_idx("out", out_idx),
+                    out[out_idx],
+                )
+    # sanity check
+    with pytest.raises(StopIteration):
+        next(out_indices)