Add handling and tests for negative and positive out-of-bounds indices

Author: amacati

src/array_api_extra/_lib/_at.py

@@ -12,9 +12,9 @@
 from ._utils._compat import (
     array_namespace,
     is_array_api_obj,
-    is_array_api_strict_namespace,
     is_dask_array,
     is_jax_array,
+    is_lazy_array,
     is_torch_array,
     is_writeable_array,
 )
@@ -390,14 +390,37 @@ def _op(
                 # We first create the mask for x
                 u_idx, _ = xp.unique_inverse(idx)
                 # Convert negative indices to positive, otherwise they won't get matched
-                u_idx = xp.where(u_idx < 0, x.shape[0] + u_idx, u_idx)
-                x_mask = xp.any(xp.arange(x.shape[0])[..., None] == u_idx, axis=-1)
+                u_idx_pos = xp.where(u_idx < 0, x.shape[0] + u_idx, u_idx)
+                # Check for out of bounds indices
+                oob_index = None
+                if is_lazy_array(u_idx_pos):
+                    pass  # Lazy arrays cannot check for out of bounds indices
+                elif xp.any(pos_oob := u_idx_pos >= x.shape[0]):
+                    first_oob_idx = xp.argmax(xp.astype(pos_oob, xp.int32))
+                    oob_index = int(u_idx[first_oob_idx])
+                elif xp.any(neg_oob := u_idx_pos < 0):
+                    first_oob_idx = xp.argmax(xp.astype(neg_oob, xp.int32))
+                    oob_index = int(u_idx[first_oob_idx])
+                if oob_index is not None:
+                    msg = (
+                        f"index {oob_index} is out of bounds for array of "
+                        f"shape {x.shape}"
+                    )
+                    raise IndexError(msg) from e
+
+                x_mask = xp.any(xp.arange(x.shape[0])[..., None] == u_idx_pos, axis=-1)
                 # If y is a scalar or 0D, we are done
                 if not is_array_api_obj(y) or y.ndim == 0:
                     x[x_mask] = y
                     return x
+                if y.shape[0] != idx.shape[0]:
+                    msg = (
+                        f"shape mismatch: value array of shape {y.shape} could not be "
+                        f"broadcast to indexing result of shape {idx.shape}"
+                    )
+                    raise ValueError(msg) from e
                 # If not, create a mask for y. Get last occurrence of each unique index
-                cmp = u_idx[:, None] == u_idx[None, :]
+                cmp = u_idx_pos[:, None] == u_idx_pos[None, :]
                 # Ignore later matches
                 lower_tri_mask = (
                     xp.arange(y.shape[0])[:, None] >= xp.arange(y.shape[0])[None, :]

tests/test_at.py

@@ -13,7 +13,9 @@
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import (
     array_namespace,
+    is_array_api_strict_namespace,
     is_jax_namespace,
+    is_numpy_namespace,
     is_writeable_array,
 )
 from array_api_extra._lib._utils._compat import device as get_device
@@ -318,6 +320,20 @@ def test_setitem_int_array_index(xp: ModuleType):
     x = xp.asarray([0.0, 1.0])
     z = at_op(x, xp.asarray([-1]), _AtOp.SET, 2.0)
     xp_assert_equal(z, xp.asarray([0.0, 2.0]))
+    # Different frameworks have all kinds of different behaviours for negative indices,
+    # out-of-bounds indices, etc. Therefore, we only test the behaviour of two
+    # frameworks: numpy because we state in the docs that it is our reference for the
+    # behaviour of other frameworks with no native support, and array-api-strict.
+    if is_array_api_strict_namespace(xp) or is_numpy_namespace(xp):
+        # Test wrong shapes
+        with pytest.raises(ValueError, match="shape"):
+            at_op(xp.asarray([0]), xp.asarray([0]), _AtOp.SET, xp.asarray([1, 2]))
+        # Test positive out of bounds index
+        with pytest.raises(IndexError, match="out of bounds"):
+            at_op(xp.asarray([0]), xp.asarray([1]), _AtOp.SET, xp.asarray([1]))
+        # Test negative out of bounds index
+        with pytest.raises(IndexError, match="out of bounds"):
+            at_op(xp.asarray([0]), xp.asarray([-2]), _AtOp.SET, xp.asarray([1]))
 
 
 @pytest.mark.parametrize("bool_mask", [False, True])