Merge pull request #385 from ev-br/pr/192_

`test_take`: test axis being optional if `x.ndim==1`

Author: ev-br

array_api_tests/test_indexing_functions.py

@@ -17,32 +17,39 @@
 )
 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")
+
+    # axis is optional but only if x.ndim == 1
+    _axis_st = st.integers(-x.ndim, max(x.ndim - 1, 0))
+    if x.ndim == 1:
+        kw = data.draw(hh.kwargs(axis=_axis_st))
+    else:
+        kw = {"axis": data.draw(_axis_st)}
+    axis = kw.get("axis", 0)
     _indices = data.draw(
         st.lists(st.integers(0, x.shape[axis] - 1), min_size=1, unique=True),
         label="_indices",
     )
+    n_axis = axis if axis>=0 else x.ndim + axis
     indices = xp.asarray(_indices, dtype=dh.default_int)
     note(f"{indices=}")
 
-    out = xp.take(x, indices, axis=axis)
+    out = xp.take(x, indices, **kw)
 
     ph.assert_dtype("take", in_dtype=x.dtype, out_dtype=out.dtype)
     ph.assert_shape(
         "take",
         out_shape=out.shape,
-        expected=x.shape[:axis] + (len(_indices),) + x.shape[axis + 1 :],
+        expected=x.shape[:n_axis] + (len(_indices),) + x.shape[n_axis + 1:],
         kw=dict(
             x=x,
             indices=indices,
             axis=axis,
         ),
     )
     out_indices = sh.ndindex(out.shape)
-    axis_indices = list(sh.axis_ndindex(x.shape, axis))
+    axis_indices = list(sh.axis_ndindex(x.shape, n_axis))
     for axis_idx in axis_indices:
         f_axis_idx = sh.fmt_idx("x", axis_idx)
         for i in _indices:
@@ -62,7 +69,6 @@ def test_take(x, data):
         next(out_indices)
 
 
-
 @pytest.mark.unvectorized
 @pytest.mark.min_version("2024.12")
 @given(