Add docs. Fix 0D and scalar cases. Handle negative indices

Author: amacati

src/array_api_extra/_lib/_at.py

@@ -150,6 +150,19 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
         >>> xpx.at(jnp.asarray([123]), jnp.asarray([0, 0])).add(1)
         Array([125], dtype=int32)
 
+    For frameworks that don't support fancy indexing by default, e.g. array-api-strict,
+    we implement a workaround for 1D integer indices and ``xpx.at().set``. Assignments
+    with multiple occurences of the same index always choose the last occurence. This is
+    consistent with numpy's behaviour, e.g.::
+
+        >>> import numpy as np
+        >>> import array_api_strict as xp
+        >>> import array_api_extra as xpx
+        >>> xpx.at(np.asarray([0]), np.asarray([0, 0])).set(np.asarray([2, 3]))
+        array([3])
+        >>> xpx.at(xp.asarray([0]), xp.asarray([0, 0])).set(xp.asarray([2, 3]))
+        Array([3], dtype=array_api_strict.int64)
+
     See Also
     --------
     jax.numpy.ndarray.at : Equivalent array method in JAX.
@@ -354,39 +367,50 @@ def _op(
         if is_torch_array(y):
             y = xp.astype(y, x.dtype, copy=False)
 
-        # Work around lack of fancy indexing __setitem__ support in array-api-strict.
-        if (
-            is_array_api_strict_namespace(xp)
-            and is_array_api_obj(idx)
-            and xp.isdtype(idx.dtype, "integral")
-            and out_of_place_op is None  # only use for set()
-        ):
-            # Vectorize the operation using boolean indexing
-            # For non-unique indices, take the last occurrence. This requires creating
-            # masks for x and y that create matching shapes.
-            unique_indices, _ = xp.unique_inverse(idx)
-            x_mask = xp.any(xp.arange(x.shape[0])[..., None] == unique_indices, axis=-1)
-            # Get last occurrence of each unique index
-            cmp = unique_indices[:, None] == unique_indices[None, :]
-            # Ignore later matches
-            lower_tri_mask = (
-                xp.arange(y.shape[0])[:, None] >= xp.arange(y.shape[0])[None, :]
-            )
-            masked_cmp = cmp & lower_tri_mask
-            # For each position i, count how many matches occurred before i
-            prior_matches = xp.sum(xp.astype(masked_cmp, xp.int32), axis=-1)
-            # Last occurrence has highest match count
-            y_mask = prior_matches == xp.max(prior_matches, axis=-1)
-            # Apply the operation only to last occurrences
-            x[x_mask] = y[y_mask]
-            return x
-
         # Backends without boolean indexing (other than JAX) crash here
         if in_place_op:  # add(), subtract(), ...
             x[idx] = in_place_op(x[idx], y)
-        else:  # set()
+            return x
+        # set()
+        try:  # We first try to use the backend's __setitem__ if available
             x[idx] = y
-        return x
+            return x
+        except IndexError as e:
+            if "Fancy indexing" not in str(e):  # Avoid masking other index errors
+                raise e
+            # Work around lack of fancy indexing __setitem__
+            if (
+                is_array_api_obj(idx)
+                and xp.isdtype(idx.dtype, "integral")
+                and idx.ndim == 1
+            ):
+                # Vectorize the operation using boolean indexing
+                # For non-unique indices, take the last occurrence. This requires
+                # masks for x and y that create matching shapes.
+                # 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)
+                # 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 not, create a mask for y. Get last occurrence of each unique index
+                cmp = u_idx[:, None] == u_idx[None, :]
+                # Ignore later matches
+                lower_tri_mask = (
+                    xp.arange(y.shape[0])[:, None] >= xp.arange(y.shape[0])[None, :]
+                )
+                masked_cmp = cmp & lower_tri_mask
+                # For each position i, count how many matches occurred before i
+                prior_matches = xp.sum(xp.astype(masked_cmp, xp.int32), axis=-1)
+                # Last occurrence has highest match count
+                y_mask = prior_matches == xp.max(prior_matches, axis=-1)
+                # Apply the operation only to last occurrences
+                x[x_mask] = y[y_mask]
+                return x
+            raise e
 
     def set(
         self,

tests/test_at.py

@@ -306,6 +306,18 @@ def test_setitem_int_array_index(xp: ModuleType):
     expect = xp.asarray([[4.0, 5.0], [2.0, 3.0]])
     z = at_op(x, idx, _AtOp.SET, y)
     xp_assert_equal(z, expect)
+    # Scalar
+    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]))
+    # 0D array
+    x = xp.asarray([0.0, 1.0])
+    z = at_op(x, xp.asarray([1]), _AtOp.SET, xp.asarray(2.0))
+    xp_assert_equal(z, xp.asarray([0.0, 2.0]))
+    # Negative indices
+    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]))
 
 
 @pytest.mark.parametrize("bool_mask", [False, True])