WIP lazywhere

Author: crusaderky

src/array_api_extra/_lib/_funcs.py

@@ -5,17 +5,18 @@
 
 import math
 import warnings
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from types import ModuleType
-from typing import cast
+from typing import cast, overload
 
 from ._at import at
 from ._utils import _compat, _helpers
-from ._utils._compat import array_namespace, is_jax_array
+from ._utils._compat import array_namespace, is_array_api_obj, is_jax_array
 from ._utils._helpers import asarrays
 from ._utils._typing import Array
 
 __all__ = [
+    "apply_where",
     "atleast_nd",
     "cov",
     "create_diagonal",
@@ -28,6 +29,128 @@
 ]
 
 
+@overload
+def apply_where(
+    cond: Array,
+    f1: Callable[..., Array],
+    f2: Callable[..., Array],
+    /,
+    *args: Array,
+    xp: ModuleType | None = None,
+): ...
+
+
+@overload
+def apply_where(
+    cond: Array,
+    f1: Callable[..., Array],
+    /,
+    *args: Array,
+    fill_value: Array | int | float | complex | bool,
+    xp: ModuleType | None = None,
+): ...
+
+
+def apply_where(
+    cond: Array,
+    f1: Callable[..., Array],
+    f2_or_args0: Callable[..., Array] | Array,
+    /,
+    *args: Array,
+    fill_value: Array | int | float | complex | bool | None = None,
+    xp: ModuleType | None = None,
+):
+    """Return elements chosen from two possibilities depending on a condition
+
+    Equivalent to ``f1(*args) if cond else f2(*args)`` performed elementwise.
+
+    Parameters
+    ----------
+    cond : array
+        The condition (expressed as a boolean array).
+    f1 : callable
+        Where `cond` is True, output will be ``f1(arr1[cond], arr2[cond], ...)``.
+    f2 : callable, optional
+        Where `cond` is False, output will be ``f1(arr1[cond], arr2[cond], ...)``.
+        Mutually exclusive with `fill_value`.
+    *args : one or more arrays
+        Arguments to `f1` (and `f2`). Must be broadcastable with `cond`.
+    fill_value : Array or scalar, optional
+        If provided, value with which to fill output array where `cond` is
+        not True. Mutually exclusive with `f2`. You must provide either one.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `cond` and `args`. Default: infer.
+
+    Returns
+    -------
+    out : array
+        An array with elements from the output of `f1` where `cond` is True and either
+        the output of `f2` or `fill_value` where `cond` is False. The returned array has
+        data type determined by Type Promotion Rules between the output of `f` and
+        either `fill_value` or the output of `f2`.
+
+    Notes
+    -----
+    ``xp.where(cond, f1(*args), f2(*args))`` requires explicitly evaluating f1 even when
+    `cond` is False, and `f2` when cond is True. This function evaluates each function
+    only for their matching condition when the backend allows for it.
+
+    Examples
+    --------
+    >>> a, b = xp.asarray([1, 2, 3, 4]), xp.asarray([5, 6, 7, 8])
+    >>> def f(a, b):
+    ...     return a * b
+    >>> apply_where(a > 2, f, a, b, fill_value=xp.nan)
+    array([ nan,  nan,  21.,  32.])
+
+    """
+    # Parse and normalize arguments
+    mutually_exc_msg = "Exactly one of `fill_value` or `f2` must be given."
+    if is_array_api_obj(f2_or_args0):
+        args = (cast(Array, f2_or_args0), *args)
+        if fill_value is not None:
+            raise TypeError(mutually_exc_msg)
+        f2: Callable[..., Array] | None = None
+    else:
+        if not callable(f2):
+            msg = "Third parameter must be either an Array or callable."
+            raise ValueError(msg)
+        f2 = f2_or_args0
+        if fill_value is None:
+            raise TypeError(mutually_exc_msg)
+
+    xp = array_namespace(cond, *args) if xp is None else xp
+
+    if fill_value is not None and getattr(fill_value, "ndim", 0) != 0:
+        msg = "`fill_value` must be a scalar."
+        raise ValueError(msg)
+
+    args = xp.broadcast_arrays(cond, *args)
+    bool_dtype = xp.asarray([True]).dtype  # numpy 1.xx doesn't have `bool`
+    cond, args = xp.astype(args[0], bool_dtype, copy=False), args[1:]
+
+    temp1 = f1(*(arr[cond] for arr in args))
+
+    if f2 is None:
+        if is_array_api_obj(fill_value) or xp.__array_api_version__ >= "2024.12":
+            dtype = xp.result_type(temp1.dtype, fill_value)
+        else:
+            # TODO: remove this branch when all backends support
+            # Array API 2024.12
+            dtype = (xp.zeros((), dtype=temp1.dtype) * fill_value).dtype
+        out = xp.full(
+            cond.shape, dtype=dtype, fill_value=xp.asarray(fill_value, dtype=dtype)
+        )
+    else:
+        ncond = ~cond
+        temp2 = xp.asarray(f2(*(arr[ncond] for arr in args)))
+        dtype = xp.result_type(temp1, temp2)
+        out = xp.empty(cond.shape, dtype=dtype)
+        out = at(out, ncond).set(temp2)
+
+    return at(out, cond).set(temp1)
+
+
 def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
     """
     Recursively expand the dimension of an array to at least `ndim`.