v2

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -7,10 +7,11 @@
 """
 from __future__ import annotations
 
+import operator
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Optional, Union, Any
+    from typing import Callable, Optional, Union, Any
     from ._typing import Array, Device
 
 import sys
@@ -811,9 +812,22 @@ def is_writeable_array(x):
         return False
     return True
 
+def _parse_copy_param(x, copy: bool | None) -> bool:
+    """Preprocess and validate a copy parameter, in line with the same
+    parameter in np.asarray(), np.astype(), etc.
+    """
+    if copy is None:
+        return not is_writeable_array(x)
+    if copy is False:
+        if not is_writeable_array(x):
+            raise ValueError("Cannot avoid modifying parameter in place")
+    elif copy is not True:
+        raise ValueError(f"Invalid value for copy: {copy!r}")
+    return copy
+
 _undef = object()
 
-def at(x, idx=_undef, /):
+class at:
     """
     Update operations for read-only arrays.
 
@@ -823,12 +837,22 @@ def at(x, idx=_undef, /):
     Keyword arguments (e.g. ``indices_are_sorted``) are passed to JAX and are
     quietly ignored for backends that don't support them.
 
+    Additionally, this introduces support for the `copy` keyword for all backends:
+
+    None
+        x *may* be modified in place if it is possible and beneficial
+        for performance. You should not use x after calling this function.
+    True
+        Ensure that the inputs are not modified. This is the default.
+    False
+        Raise ValueError if a copy cannot be avoided.
+
     Examples
     --------
     Given either of these equivalent expressions::
 
-      x = at(x)[1].add(2)
-      x = at(x, 1).add(2)
+      x = at(x)[1].add(2, copy=None)
+      x = at(x, 1).add(2, copy=None)
 
     If x is a JAX array, they are the same as::
 
@@ -845,16 +869,17 @@ def at(x, idx=_undef, /):
 
     Warning
     -------
-    You should always immediately overwrite the parameter array::
+    When you use copy=None, you should always immediately overwrite
+    the parameter array::
 
-        x = at(x, 0).set(2)
+        x = at(x, 0).set(2, copy=None)
 
     The anti-pattern below must be avoided, as it will result in different behaviour
     on read-only versus writeable arrays:
 
         x = xp.asarray([0, 0, 0])
-        y = at(x, 0).set(2)
-        z = at(x, 1).set(3)
+        y = at(x, 0).set(2, copy=None)
+        z = at(x, 1).set(3, copy=None)
 
     In the above example, y == [2, 0, 0] and z == [0, 3, 0] when x is read-only,
     whereas y == z == [2, 3, 0] when x is writeable!
@@ -863,18 +888,6 @@ def at(x, idx=_undef, /):
     --------
     https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html
     """
-    if is_jax_array(x):
-        return x.at
-    if is_numpy_array(x) and not x.flags.writeable:
-        x = x.copy()
-    return _InPlaceAt(x, idx)
-
-class _InPlaceAt:
-    """Helper of at().
-
-    Trivially implement jax.numpy.ndarray.at for other backends.
-    x is updated in place.
-    """
     __slots__ = ("x", "idx")
 
     def __init__(self, x, idx=_undef):
@@ -890,7 +903,16 @@ def __getitem__(self, idx):
         self.idx = idx
         return self
 
-    def _check_args(self, mode="promise_in_bounds", **kwargs):
+    def _common(self, at_op, y=_undef, mode: str = "promise_in_bounds", **kwargs):
+        """Validate kwargs and perform common prepocessing.
+
+        Returns
+        -------
+        If the operation can be resolved by at[],
+            (return value, None)
+        Otherwise,
+            (None, preprocessed x)
+        """
         if self.idx is _undef:
             raise TypeError(
                 "Index has not been set.\n"
@@ -900,74 +922,136 @@ def _check_args(self, mode="promise_in_bounds", **kwargs):
                 "    at(x)[idx].set(value)\n"
                 "(same for all other methods)."
             )
-        if mode != "promise_in_bounds":
+        if mode != "promise_in_bounds" and not is_jax_array(self.x):
             xp = array_namespace(self.x)
             raise NotImplementedError(
-                f"mode='{mode}' is not supported for backend {xp.__name__}"
+                f"mode='{mode!r}' is not supported for backend {xp.__name__}"
+            )
+
+        copy = _parse_copy_param(self.x, copy)
+
+        if copy and is_jax_array(self.x):
+            # Use JAX's at[]
+            at_ = self.x.at[self.idx]
+            args = (y, ) if y is not _undef else ()
+            return getattr(at_, at_op)(*args, mode=mode, **kwargs), None
+
+        # Emulate at[] behaviour for non-JAX arrays
+        x = self.x.copy() if copy else self.x
+        return None, x
+
+    def get(self, copy: bool | None = True, **kwargs):
+        """Return x[idx]. In addition to plain __getitem__, this allows ensuring
+        that the output is (not) a copy and kwargs are passed to the backend."""
+        # Special case when xp=numpy and idx is a fancy index
+        # If copy is not False, avoid an unnecessary double copy.
+        # if copy is forced to False, raise.
+        if (
+            is_numpy_array(self.x) 
+            and (
+                isinstance(self.idx, (list, tuple))
+                or (is_numpy_array(self.idx) and self.idx.dtype.kind in "biu")
             )
+        ):
+            if copy is True:
+                copy = None
+            elif copy is False:
+                raise ValueError(
+                    "Indexing a numpy array with a fancy index always "
+                    "results in a copy"
+                )
+
+        res, x = self._common("get", copy=copy, **kwargs)
+        if res is not None:
+            return res
+        return x[self.idx]
 
     def set(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] = y
-        return self.x
+        """x[idx] = y"""
+        res, x = self._common("set", y, **kwargs)
+        if res is not None:
+            return res
+        x[self.idx] = y
+        return x
+
+    def apply(self, ufunc, /, **kwargs):
+        """ufunc.at(x, idx)"""
+        res, x = self._common("apply", ufunc, **kwargs)
+        if res is not None:
+            return res
+        ufunc.at(x, self.idx)
+        return x
+
+    def _iop(self, at_op: str, elwise_op: Callable[[Array, Array], Array], y: Array, **kwargs):
+        """x[idx] += y or equivalent in-place operation on a subset of x
+
+        which is the same as saying
+            x[idx] = x[idx] + y
+        Note that this is not the same as 
+            operator.iadd(x[idx], y)
+        Consider for example when x is a numpy array and idx is a fancy index, which
+        triggers a deep copy on __getitem__.
+        """
+        res, x = self._common(at_op, y, **kwargs)
+        if res is not None:
+            return res
+        x[self.idx] = elwise_op(x[self.idx], y)
+        return x
 
     def add(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] += y
-        return self.x
-    
+        """x[idx] += y"""
+        return self._iop("add", operator.add, y, **kwargs)
+
     def subtract(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] -= y
-        return self.x
+        """x[idx] -= y"""
+        return self._iop("subtract", operator.sub, y, **kwargs)
 
     def multiply(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] *= y
-        return self.x
+        """x[idx] *= y"""
+        return self._iop("multiply", operator.mul, y, **kwargs)
 
     def divide(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] /= y
-        return self.x
-    
+        """x[idx] /= y"""
+        return self._iop("divide", operator.truediv, y, **kwargs)
+
     def power(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        self.x[self.idx] **= y
-        return self.x
+        """x[idx] **= y"""
+        return self._iop("power", operator.pow, y, **kwargs)
 
     def min(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        xp = array_namespace(self.x, y)
-        self.x[self.idx] = xp.minimum(self.x[self.idx], y)
-        return self.x
+        """x[idx] = minimum(x[idx], y)"""
+        xp = array_namespace(self.x)
+        return self._iop("min", xp.minimum, y, **kwargs)
 
     def max(self, y, /, **kwargs):
-        self._check_args(**kwargs)
-        xp = array_namespace(self.x, y)
-        self.x[self.idx] = xp.maximum(self.x[self.idx], y)
-        return self.x
-
-    def apply(self, ufunc, /, **kwargs):
-        self._check_args(**kwargs)
-        ufunc.at(self.x, self.idx)
-        return self.x
-
-    def get(self, **kwargs):
-        self._check_args(**kwargs)
-        return self.x[self.idx]
-
-def iwhere(condition, x, y, /):
-    """Variant of xp.where(condition, x, y) which may or may not update
-    x in place, if it's possible and beneficial for performance.    
+        """x[idx] = maximum(x[idx], y)"""
+        xp = array_namespace(self.x)
+        return self._iop("max", xp.maximum, y, **kwargs)
+
+def where(condition, x, y, /, copy: bool | None = True):
+    """Return elements from x when condition is True and from y when
+    it is False. 
+
+    This is a wrapper around xp.where that adds the copy parameter:
+
+    None
+        x *may* be modified in place if it is possible and beneficial
+        for performance. You should not use x after calling this function.
+    True
+        Ensure that the inputs are not modified.
+        This is the default, in line with np.where.
+    False
+        Raise ValueError if a copy cannot be avoided.
     """
+    copy = _parse_copy_param(x, copy)
     xp = array_namespace(condition, x, y)
-    if is_writeable_array(x):
+    if copy:
+        return xp.where(condition, x, y)
+    else:
         condition, x, y = xp.broadcast_arrays(condition, x, y)
         x[condition] = y[condition]
         return x
-    else:
-        return xp.where(condition, x, y)
+
 
 __all__ = [
     "array_namespace",
@@ -993,7 +1077,7 @@ def iwhere(condition, x, y, /):
     "size",
     "to_device",
     "at",
-    "iwhere",
+    "where",
 ]
 
 _all_ignore = ['sys', 'math', 'inspect', 'warnings']

docs/helper-functions.rst

@@ -37,7 +37,7 @@ instead, which would be wrapped.
 .. autofunction:: to_device
 .. autofunction:: size
 .. autofunction:: at
-.. autofunction:: iwhere
+.. autofunction:: where
 
 Inspection Helpers
 ------------------