ENH: allow ndim>1 indexing arrays

Author: ev-br

array_api_strict/_array_object.py

@@ -498,14 +498,9 @@ def _validate_index(self, key, op="getitem"):
                             "Array API when the array is the sole index."
                         )
                     if not get_array_api_strict_flags()['boolean_indexing']:
-                        raise RuntimeError("The boolean_indexing flag has been disabled for array-api-strict")
-
-                elif i.dtype in _integer_dtypes and i.ndim > 1:
-                    raise IndexError(
-                        f"Single-axes index {i} is a multi-dimensional "
-                        "integer array, but advanced integer indexing is only "
-                        "specified in the Array API for 1D index arrays."
-                    )
+                        raise RuntimeError(
+                            "The boolean_indexing flag has been disabled for array-api-strict"
+                        )
             elif isinstance(i, tuple):
                 raise IndexError(
                     f"Single-axes index {i} is a tuple, but nested tuple "

array_api_strict/tests/test_array_object.py

@@ -70,11 +70,6 @@ def test_validate_index():
     assert_raises(IndexError, lambda: a[[True, True, True]])
     assert_raises(IndexError, lambda: a[(True, True, True),])
 
-    # Integer array indices are not allowed (except for 0-D or 1D)
-    idx = asarray([[0, 1]])  # idx.ndim == 2
-    assert_raises(IndexError, lambda: a[idx, 0])
-    assert_raises(IndexError, lambda: a[0, idx])
-
     # Mixing 1D integer array indices with slices, ellipsis or booleans is not allowed
     idx = asarray([0, 1])
     assert_raises(IndexError, lambda: a[..., idx])
@@ -138,6 +133,10 @@ def test_indexing_arrays():
     with assert_raises(IndexError):
         a[idx, idx] = 42
 
+    # smoke test indexing with ndim > 1 arrays
+    idx = idx[..., None]
+    a[idx, idx]
+
 
 def test_promoted_scalar_inherits_device():
     device1 = Device("device1")