ENH: allow 1D integer array indices

Author: ev-br

array_api_strict/_array_object.py

@@ -395,6 +395,8 @@ def _validate_index(self, key):
         single_axes = []
         n_ellipsis = 0
         key_has_mask = False
+        key_has_index_array = False
+        key_has_slices = False
         for i in _key:
             if i is not None:
                 nonexpanding_key.append(i)
@@ -403,13 +405,17 @@ def _validate_index(self, key):
                 if isinstance(i, Array):
                     if i.dtype in _boolean_dtypes:
                         key_has_mask = True
+                    elif i.dtype in _integer_dtypes:
+                        key_has_index_array = True
                     single_axes.append(i)
                 else:
                     # i must not be an array here, to avoid elementwise equals
                     if i == Ellipsis:
                         n_ellipsis += 1
                     else:
                         single_axes.append(i)
+                        if isinstance(i, slice):
+                            key_has_slices = True
 
         n_single_axes = len(single_axes)
         if n_ellipsis > 1:
@@ -427,6 +433,12 @@ def _validate_index(self, key):
                     "specified in the Array API."
                 )
 
+        if (key_has_index_array and (n_ellipsis > 0 or key_has_slices or key_has_mask)):
+            raise IndexError(
+                "Integer index arrays are only allowed with integer indices; "
+                f"got {key}."
+            )
+
         if n_ellipsis == 0:
             indexed_shape = self.shape
         else:
@@ -485,11 +497,11 @@ def _validate_index(self, key):
                     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 != 0:
+                elif i.dtype in _integer_dtypes and i.ndim > 1:
                     raise IndexError(
-                        f"Single-axes index {i} is a non-zero-dimensional "
-                        "integer array, but advanced integer indexing is not "
-                        "specified in the Array API."
+                        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."
                     )
             elif isinstance(i, tuple):
                 raise IndexError(

array_api_strict/tests/test_array_object.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-from .. import ones, asarray, result_type, all, equal
+from .. import ones, arange, reshape, asarray, result_type, all, equal
 from .._array_object import Array, CPU_DEVICE, Device
 from .._dtypes import (
     _all_dtypes,
@@ -70,11 +70,25 @@ 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)
-    idx = asarray([0, 1])
+    # 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])
+    assert_raises(IndexError, lambda: a[:, idx])
+    assert_raises(IndexError, lambda: a[asarray([True, True]), idx])
+
+    # 1D integer array indices must have the same length
+    idx1 = asarray([0, 1])
+    idx2 = asarray([0, 1, 1])
+    assert_raises(IndexError, lambda: a[idx1, idx2])
+
+    # Non-integer array indices are not allowed
+    assert_raises(IndexError, lambda: a[ones(2), 0])
+
     # Array-likes (lists, tuples) are not allowed as indices
     assert_raises(IndexError, lambda: a[[0, 1]])
     assert_raises(IndexError, lambda: a[(0, 1), (0, 1)])
@@ -91,6 +105,37 @@ def test_validate_index():
     assert_raises(IndexError, lambda: a[:])
     assert_raises(IndexError, lambda: a[idx])
 
+
+def test_indexing_arrays():
+    # indexing with 1D integer arrays and mixes of integers and 1D integer are allowed
+
+    # 1D array
+    a = arange(5)
+    idx = asarray([1, 0, 1, 2, -1])
+    a_idx = a[idx]
+
+    a_idx_loop = asarray([a[idx[i]] for i in range(idx.shape[0])])
+    assert all(a_idx == a_idx_loop)
+
+    # setitem with arrays is not allowed      # XXX
+    # with assert_raises(IndexError):
+    #    a[idx] = 42
+
+    # mixed array and integer indexing
+    a = reshape(arange(3*4), (3, 4))
+    idx = asarray([1, 0, 1, 2, -1])
+    a_idx = a[idx, 1]
+
+    a_idx_loop = asarray([a[idx[i], 1] for i in range(idx.shape[0])])
+    assert all(a_idx == a_idx_loop)
+
+
+    # index with two arrays
+    a_idx = a[idx, idx]
+    a_idx_loop = asarray([a[idx[i], idx[i]] for i in range(idx.shape[0])])
+    assert all(a_idx == a_idx_loop)
+
+
 def test_promoted_scalar_inherits_device():
     device1 = Device("device1")
     x = asarray([1., 2, 3], device=device1)