ENH: only test with 1D index arrays, make test unvecorized

Author: ev-br

array_api_tests/test_array_object.py

@@ -242,8 +242,25 @@ def test_setitem_masking(shape, data):
             )
 
 
-@given(shape=hh.shapes(), data=st.data())
-def test_getitem_arrays_and_ints(shape, data):
+@pytest.mark.min_version("2024.12")
+@pytest.mark.unvectorized
+@given(shape=hh.shapes(min_dims=2), data=st.data())
+def test_getitem_arrays_and_ints_1(shape, data):
+    # min_dims=2 : test multidim `x` arrays
+    # index arrays are all 1D
+    _test_getitem_arrays_and_ints_1D(shape, data)
+
+
+@pytest.mark.min_version("2024.12")
+@pytest.mark.unvectorized
+@given(shape=hh.shapes(min_dims=1), data=st.data())
+def test_getitem_arrays_and_ints_2(shape, data):
+    # min_dims=1 : favor 1D `x` arrays
+    # index arrays are all 1D
+    _test_getitem_arrays_and_ints_1D(shape, data)
+
+
+def _test_getitem_arrays_and_ints_1D(shape, data):
     assume((len(shape) > 0) and all(sh > 0 for sh in shape))
 
     dtype = xp.int32
@@ -254,10 +271,12 @@ def test_getitem_arrays_and_ints(shape, data):
     arr_index = [data.draw(st.booleans()) for _ in range(len(shape))]
     assume(sum(arr_index) > 0)
 
-    # draw shapes for index arrays
+    # draw shapes for index arrays: NB max_dims=1 ==> 1D indexing arrays ONLY
     if sum(arr_index) > 0:
         index_shapes = data.draw(
-            hh.mutually_broadcastable_shapes(sum(arr_index), min_dims=1, min_side=1)
+            hh.mutually_broadcastable_shapes(
+                sum(arr_index), min_dims=1, max_dims=1, min_side=1
+            )
         )
         index_shapes = list(index_shapes)
 
@@ -279,19 +298,18 @@ def test_getitem_arrays_and_ints(shape, data):
             # draw an integer
             key.append(data.draw(st.integers(-shape[i], shape[i]-1)))
 
-
-    print(f"??? {x.shape = } {key = } -- {[k if isinstance(k, int) else k.shape for k in key]}")
+    # print(f"??? {x.shape = } {key = }")
 
     key = tuple(key)
     out = x[key]
 
-    # XXX: how to properly check
-    import numpy as np
-    x_np = np.asarray(x)
-    out_np = np.asarray(out)
-    key_np = tuple(k if isinstance(k, int) else np.asarray(k) for k in key)
+    arrays = [xp.asarray(k) for k in key]
+    bcast_shape = sh.broadcast_shapes(*[arr.shape for arr in arrays])
+    bcast_key = [xp.broadcast_to(arr, bcast_shape) for arr in arrays]
 
-    np.testing.assert_equal(out_np, x_np[key_np])
+    for idx in sh.ndindex(bcast_shape):
+        tpl = tuple(k[idx] for k in bcast_key)
+        assert out[idx] == x[tpl], f"failing at {idx = } w/ {key = }"
 
 
 def make_scalar_casting_param(