ENH: test with ND index arrays

Author: ev-br

array_api_tests/test_array_object.py

@@ -242,25 +242,29 @@ def test_setitem_masking(shape, data):
             )
 
 
+# ### Fancy indexing ###
+
 @pytest.mark.min_version("2024.12")
 @pytest.mark.unvectorized
+@pytest.mark.parametrize("idx_max_dims", [1, None])
 @given(shape=hh.shapes(min_dims=2), data=st.data())
-def test_getitem_arrays_and_ints_1(shape, data):
+def test_getitem_arrays_and_ints_1(shape, data, idx_max_dims):
     # min_dims=2 : test multidim `x` arrays
-    # index arrays are all 1D
-    _test_getitem_arrays_and_ints_1D(shape, data)
+    # index arrays are 1D for idx_max_dims=1 and multidim for idx_max_dims=None
+    _test_getitem_arrays_and_ints(shape, data, idx_max_dims)
 
 
 @pytest.mark.min_version("2024.12")
 @pytest.mark.unvectorized
+@pytest.mark.parametrize("idx_max_dims", [1, None])
 @given(shape=hh.shapes(min_dims=1), data=st.data())
-def test_getitem_arrays_and_ints_2(shape, data):
+def test_getitem_arrays_and_ints_2(shape, data, idx_max_dims):
     # min_dims=1 : favor 1D `x` arrays
-    # index arrays are all 1D
-    _test_getitem_arrays_and_ints_1D(shape, data)
+    # index arrays are 1D for idx_max_dims=1 and multidim for idx_max_dims=None
+    _test_getitem_arrays_and_ints(shape, data, idx_max_dims)
 
 
-def _test_getitem_arrays_and_ints_1D(shape, data):
+def _test_getitem_arrays_and_ints(shape, data, idx_max_dims):
     assume((len(shape) > 0) and all(sh > 0 for sh in shape))
 
     dtype = xp.int32
@@ -271,11 +275,12 @@ def _test_getitem_arrays_and_ints_1D(shape, data):
     arr_index = [data.draw(st.booleans()) for _ in range(len(shape))]
     assume(sum(arr_index) > 0)
 
-    # draw shapes for index arrays: NB max_dims=1 ==> 1D indexing arrays ONLY
+    # draw shapes for index arrays: max_dims=1 ==> 1D indexing arrays ONLY
+    #                               max_dims=None ==> multidim indexing arrays
     if sum(arr_index) > 0:
         index_shapes = data.draw(
             hh.mutually_broadcastable_shapes(
-                sum(arr_index), min_dims=1, max_dims=1, min_side=1
+                sum(arr_index), min_dims=1, max_dims=idx_max_dims, min_side=1
             )
         )
         index_shapes = list(index_shapes)
@@ -298,7 +303,7 @@ def _test_getitem_arrays_and_ints_1D(shape, data):
             # draw an integer
             key.append(data.draw(st.integers(-shape[i], shape[i]-1)))
 
-    # print(f"??? {x.shape = } {key = }")
+    print(f"??? {x.shape = } {len(key) = }  {[xp.asarray(k).shape for k in key]}")
 
     key = tuple(key)
     out = x[key]