ENH: apply_where (from scipy's lazywhere)

Author: crusaderky

docs/api-reference.md

@@ -6,6 +6,7 @@
     :nosignatures:
     :toctree: generated
 
+    apply_where
     at
     atleast_nd
     broadcast_shapes

pixi.lock

@@ -191,10 +191,6 @@ markers = [
 
 [tool.coverage]
 run.source = ["array_api_extra"]
-report.exclude_also = [
-    '\.\.\.',
-    'if TYPE_CHECKING:',
-]
 
 # mypy
 
@@ -236,6 +232,10 @@ reportMissingTypeStubs = false
 reportUnreachable = false
 # ruff handles this
 reportUnusedParameter = false
+# cyclic imports inside function bodies
+reportImportCycles = false
+# PyRight can't trace types in lambdas
+reportUnknownLambdaType = false
 
 executionEnvironments = [
     { root = "tests", reportPrivateUsage = false },
@@ -301,6 +301,7 @@ messages_control.disable = [
     "missing-function-docstring", # numpydoc handles this
     "import-error",               # mypy handles this
     "import-outside-toplevel",    # optional dependencies
+    "cyclic-import",              # cyclic imports inside function bodies
 ]
 
 

pyproject.toml

@@ -3,6 +3,7 @@
 from ._delegation import isclose, pad
 from ._lib._at import at
 from ._lib._funcs import (
+    apply_where,
     atleast_nd,
     broadcast_shapes,
     cov,
@@ -19,6 +20,7 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "apply_where",
     "at",
     "atleast_nd",
     "broadcast_shapes",

src/array_api_extra/__init__.py

@@ -249,14 +249,16 @@ def _op(
             Right-hand side of the operation.
         copy : bool or None
             Whether to copy the input array. See the class docstring for details.
-        xp : array_namespace or None
-            The array namespace for the input array.
+        xp : array_namespace, optional
+            The array namespace for the input array. Default: infer.
 
         Returns
         -------
         Array
             Updated `x`.
         """
+        from ._funcs import apply_where
+
         x, idx = self._x, self._idx
         xp = array_namespace(x, y) if xp is None else xp
 
@@ -294,8 +296,10 @@ def _op(
             y_xp = xp.asarray(y, dtype=x.dtype)
             if y_xp.ndim == 0:
                 if out_of_place_op:
-                    # FIXME: suppress inf warnings on dask with lazywhere
-                    out = xp.where(idx, out_of_place_op(x, y_xp), x)
+                    # suppress inf warnings on Dask
+                    out = apply_where(
+                        idx, out_of_place_op, (x, y_xp), fill_value=x, xp=xp
+                    )
                     # Undo int->float promotion on JAX after _AtOp.DIVIDE
                     out = xp.astype(out, x.dtype, copy=False)
                 else:

src/array_api_extra/_lib/_at.py

@@ -5,17 +5,23 @@
 
 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._helpers import asarrays
+from ._utils._compat import (
+    array_namespace,
+    is_dask_namespace,
+    is_jax_array,
+    is_jax_namespace,
+)
+from ._utils._helpers import asarrays, meta_namespace
 from ._utils._typing import Array
 
 __all__ = [
+    "apply_where",
     "atleast_nd",
     "broadcast_shapes",
     "cov",
@@ -29,6 +35,156 @@
 ]
 
 
+@overload
+def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
+    cond: Array,
+    f1: Callable[..., Array],
+    f2: Callable[..., Array],
+    args: Array | tuple[Array, ...],
+    /,
+    *,
+    xp: ModuleType | None = None,
+) -> Array: ...
+
+
+@overload
+def apply_where(  # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08
+    cond: Array,
+    f1: Callable[..., Array],
+    args: Array | tuple[Array, ...],
+    /,
+    *,
+    fill_value: Array | int | float | complex | bool,
+    xp: ModuleType | None = None,
+) -> Array: ...
+
+
+def apply_where(  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02
+    cond: Array,
+    f1: Callable[..., Array],
+    f2: Callable[..., Array] | Array | tuple[Array],  # optional positional argument
+    args: Array | tuple[Array, ...] | None = None,
+    /,
+    *,
+    fill_value: Array | int | float | complex | bool | None = None,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Run one of two elementwise functions depending on a condition.
+
+    Equivalent to ``f1(*args) if cond else fill_value`` performed elementwise
+    when `fill_value` is defined, otherwise to ``f1(*args) if cond else f2(*args)``.
+
+    Parameters
+    ----------
+    cond : array
+        The condition, expressed as a boolean array.
+    f1 : callable
+        Elementwise function of `args`, returning a single array.
+        Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``.
+    f2 : callable, optional
+        Elementwise function of `args`, returning a single array.
+        Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``.
+        Mutually exclusive with `fill_value`.
+    args : Array or tuple of Arrays
+        Argument(s) 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 False.
+        It does not need to be scalar.
+        Mutually exclusive with `f2`. You must provide one or the other.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `cond` and `args`. Default: infer.
+
+    Returns
+    -------
+    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 `f1` 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, if the backend allows for it.
+
+    On Dask, f1 and f2 are applied to the individual chunks.
+
+    Examples
+    --------
+    >>> a = xp.asarray([5, 4, 3])
+    >>> b = xp.asarray([0, 2, 2])
+    >>> def f(a, b):
+    ...     return a // b
+    >>> apply_where(b != 0, f, (a, b), fill_value=xp.nan)
+    array([ nan,  2., 1.])
+    """
+    # Parse and normalize arguments
+    mutually_exc_msg = "Exactly one of `fill_value` or `f2` must be given."
+    if args is None:
+        f2, args = None, f2
+        if fill_value is None:
+            raise TypeError(mutually_exc_msg)
+    else:
+        if not callable(f2):
+            msg = "Third parameter must be a callable, Array, or tuple of Arrays."
+            raise TypeError(msg)
+        if fill_value is not None:
+            raise TypeError(mutually_exc_msg)
+
+    if not isinstance(args, tuple):
+        args = (args,)
+    f2 = cast(Callable[..., Array] | None, f2)  # type: ignore[no-any-explicit]
+    args = cast(tuple[Array, ...], args)
+
+    xp = array_namespace(cond, *args) if xp is None else xp
+
+    if getattr(fill_value, "ndim", 0):
+        cond, fill_value, *args = xp.broadcast_arrays(cond, fill_value, *args)
+    else:
+        cond, *args = xp.broadcast_arrays(cond, *args)
+
+    if is_dask_namespace(xp):
+        meta_xp = meta_namespace(cond, fill_value, *args, xp=xp)
+        return xp.map_blocks(_apply_where, cond, f1, f2, fill_value, *args, xp=meta_xp)
+    return _apply_where(cond, f1, f2, fill_value, *args, xp=xp)
+
+
+def _apply_where(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01
+    cond: Array,
+    f1: Callable[..., Array],
+    f2: Callable[..., Array] | None,
+    fill_value: Array | int | float | complex | bool | None,
+    *args: Array,
+    xp: ModuleType,
+) -> Array:
+    """Helper of `apply_where`. On Dask, this runs on a single chunk."""
+
+    if is_jax_namespace(xp):
+        # jax.jit does not support assignment by boolean mask
+        return xp.where(cond, f1(*args), f2(*args) if f2 is not None else fill_value)
+
+    device = _compat.device(cond)
+    temp1 = f1(*(arr[cond] for arr in args))
+
+    if f2 is None:
+        # TODO remove asarrays once all backends support Array API 2024.12
+        dtype = xp.result_type(*asarrays(temp1, fill_value, xp=xp))
+        if getattr(fill_value, "ndim", 0):
+            fill_value = xp.astype(fill_value, dtype)
+            return at(fill_value, cond).set(temp1, copy=True)
+        out = xp.full(cond.shape, fill_value=fill_value, dtype=dtype, device=device)
+    else:
+        ncond = ~cond
+        temp2 = f2(*(arr[ncond] for arr in args))
+        dtype = xp.result_type(temp1, temp2)
+        out = xp.empty(cond.shape, dtype=dtype, device=device)
+        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`.
@@ -385,12 +541,15 @@ def isclose(
     a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating"))
     b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating"))
     if a_inexact or b_inexact:
-        # FIXME: use scipy's lazywhere to suppress warnings on inf
-        out = xp.where(
+        # prevent warnings on numpy and dask on inf - inf
+        mxp = meta_namespace(a, b, xp=xp)
+        out = apply_where(
             xp.isinf(a) | xp.isinf(b),
-            xp.isinf(a) & xp.isinf(b) & (xp.sign(a) == xp.sign(b)),
+            lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)),  # pyright: ignore[reportUnknownArgumentType]
             # Note: inf <= inf is True!
-            xp.abs(a - b) <= (atol + rtol * xp.abs(b)),
+            lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)),  # pyright: ignore[reportUnknownArgumentType]
+            (a, b),
+            xp=xp,
         )
         if equal_nan:
             out = xp.where(xp.isnan(a) & xp.isnan(b), xp.asarray(True), out)

src/array_api_extra/_lib/_funcs.py

@@ -7,10 +7,15 @@
 from typing import cast
 
 from . import _compat
-from ._compat import is_array_api_obj, is_numpy_array
+from ._compat import (
+    array_namespace,
+    is_array_api_obj,
+    is_dask_namespace,
+    is_numpy_array,
+)
 from ._typing import Array
 
-__all__ = ["in1d", "mean"]
+__all__ = ["asarrays", "in1d", "is_python_scalar", "mean", "meta_namespace"]
 
 
 def in1d(
@@ -33,7 +38,7 @@ def in1d(
     https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
     """
     if xp is None:
-        xp = _compat.array_namespace(x1, x2)
+        xp = array_namespace(x1, x2)
 
     # This code is run to make the code significantly faster
     if x2.shape[0] < 10 * x1.shape[0] ** 0.145:
@@ -84,7 +89,7 @@ def mean(
     Complex mean, https://github.com/data-apis/array-api/issues/846.
     """
     if xp is None:
-        xp = _compat.array_namespace(x)
+        xp = array_namespace(x)
 
     if xp.isdtype(x.dtype, "complex floating"):
         x_real = xp.real(x)
@@ -124,8 +129,8 @@ def asarrays(
     ----------
     a, b : Array | int | float | complex | bool
         Input arrays or scalars. At least one must be an array.
-    xp : ModuleType
-        The standard-compatible namespace for the returned arrays.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
 
     Returns
     -------
@@ -175,3 +180,33 @@ def asarrays(
         xa, xb = xp.asarray(a), xp.asarray(b)
 
     return (xb, xa) if swap else (xa, xb)
+
+
+def meta_namespace(
+    *arrays: Array | int | float | complex | bool | None,
+    xp: ModuleType | None = None,
+) -> ModuleType:
+    """
+    Get the namespace of Dask chunks.
+
+    On all other backends, just return the namespace of the arrays.
+
+    Parameters
+    ----------
+    *arrays : Array | int | float | complex | bool | None
+        Input arrays.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array_namespace
+        If xp is Dask, the namespace of the Dask chunks;
+        otherwise, the namespace of the arrays.
+    """
+    xp = array_namespace(*arrays) if xp is None else xp
+    if not is_dask_namespace(xp):
+        return xp
+    # Quietly skip scalars and None's
+    metas = [getattr(a, "_meta", None) for a in arrays]
+    return array_namespace(*metas)

src/array_api_extra/_lib/_utils/_helpers.py

@@ -4,7 +4,8 @@
 
 # To be changed to a Protocol later (see data-apis/array-api#589)
 Array = Any  # type: ignore[no-any-explicit]
+DType = Any  # type: ignore[no-any-explicit]
 Device = Any  # type: ignore[no-any-explicit]
 Index = Any  # type: ignore[no-any-explicit]
 
-__all__ = ["Array", "Device", "Index"]
+__all__ = ["Array", "DType", "Device", "Index"]

src/array_api_extra/_lib/_utils/_typing.py

@@ -18,7 +18,7 @@
 
 __all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     # TODO move ParamSpec outside TYPE_CHECKING
     # depends on scikit-learn abandoning Python 3.9
     # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
@@ -169,7 +169,7 @@ def xp(request, monkeypatch):
         Pytest fixture, as acquired by the test itself or by one of its fixtures.
     monkeypatch : pytest.MonkeyPatch
         Pytest fixture, as acquired by the test itself or by one of its fixtures.
-    xp : module
+    xp : array_namespace
         Array namespace to be tested.
 
     See Also

src/array_api_extra/testing.py

@@ -115,7 +115,11 @@ def xp(
     if library == Backend.NUMPY_READONLY:
         return NumPyReadOnly()  # type: ignore[return-value]  # pyright: ignore[reportReturnType]
     xp = pytest.importorskip(library.value)
+    # Possibly wrap module with array_api_compat
+    xp = array_namespace(xp.empty(0))
 
+    # On Dask and JAX, monkey-patch all functions tagged by `lazy_xp_function`
+    # in the global scope of the module containing the test function.
     patch_lazy_xp_functions(request, monkeypatch, xp=xp)
 
     if library == Backend.JAX:
@@ -124,8 +128,18 @@ def xp(
         # suppress unused-ignore to run mypy in -e lint as well as -e dev
         jax.config.update("jax_enable_x64", True)  # type: ignore[no-untyped-call,unused-ignore]
 
-    # Possibly wrap module with array_api_compat
-    return array_namespace(xp.empty(0))
+    return xp
+
+
+@pytest.fixture(params=[Backend.DASK])  # Can select the test with `pytest -k dask`
+def da(
+    request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> ModuleType:  # numpydoc ignore=PR01,RT01
+    """Variant of the `xp` fixture that only yields dask.array."""
+    xp = pytest.importorskip("dask.array")
+    xp = array_namespace(xp.empty(0))
+    patch_lazy_xp_functions(request, monkeypatch, xp=xp)
+    return xp
 
 
 @pytest.fixture

tests/conftest.py

@@ -241,18 +241,7 @@ def test_bool_mask_nd(xp: ModuleType):
     xp_assert_equal(z, xp.asarray([[0, 2, 3], [4, 0, 0]]))
 
 
-@pytest.mark.parametrize(
-    "bool_mask",
-    [
-        False,
-        pytest.param(
-            True,
-            marks=pytest.mark.xfail_xp_backend(
-                Backend.DASK, reason="FIXME need scipy's lazywhere"
-            ),
-        ),
-    ],
-)
+@pytest.mark.parametrize("bool_mask", [False, True])
 def test_no_inf_warnings(xp: ModuleType, bool_mask: bool):
     x = xp.asarray([math.inf, 1.0, 2.0])
     idx = ~xp.isinf(x) if bool_mask else slice(1, None)

tests/test_at.py

@@ -1,12 +1,14 @@
 import contextlib
 import math
 import warnings
+from collections.abc import Callable
 from types import ModuleType
 
 import numpy as np
 import pytest
 
 from array_api_extra import (
+    apply_where,
     at,
     atleast_nd,
     broadcast_shapes,
@@ -43,6 +45,179 @@
 lazy_xp_function(sinc, jax_jit=False, static_argnames="xp")
 
 
+def apply_where_jit(  # type: ignore[no-any-explicit]
+    cond: Array,
+    f1: Callable[..., Array],
+    f2: Callable[..., Array] | None,
+    args: Array | tuple[Array, ...],
+    fill_value: Array | int | float | complex | bool | None = None,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Work around jax.jit's inability to handle variadic positional arguments.
+
+    This is a lazy_xp_function artefact for when jax.jit is applied directly
+    to apply_where, which would not happen in real life.
+    """
+    if f2 is None:
+        return apply_where(cond, f1, args, fill_value=fill_value, xp=xp)
+    assert fill_value is None
+    return apply_where(cond, f1, f2, args, xp=xp)
+
+
+lazy_xp_function(apply_where_jit, static_argnames=("f1", "f2", "xp"))
+
+
+class TestApplyWhere:
+    @staticmethod
+    def f1(x: Array, y: Array | int = 10) -> Array:
+        return x + y
+
+    @staticmethod
+    def f2(x: Array, y: Array | int = 10) -> Array:
+        return x - y
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_f1_f2(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4])
+        cond = x % 2 == 0
+        actual = apply_where_jit(cond, self.f1, self.f2, x)
+        expect = xp.where(cond, self.f1(x), self.f2(x))
+        xp_assert_equal(actual, expect)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_fill_value(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4])
+        cond = x % 2 == 0
+        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=0)
+        expect = xp.where(cond, self.f1(x), xp.asarray(0))
+        xp_assert_equal(actual, expect)
+
+        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=xp.asarray(0))
+        xp_assert_equal(actual, expect)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_args_tuple(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4])
+        y = xp.asarray([10, 20, 30, 40])
+        cond = x % 2 == 0
+        actual = apply_where_jit(cond, self.f1, self.f2, (x, y))
+        expect = xp.where(cond, self.f1(x, y), self.f2(x, y))
+        xp_assert_equal(actual, expect)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_broadcast(self, xp: ModuleType):
+        x = xp.asarray([1, 2])
+        y = xp.asarray([[10], [20], [30]])
+        cond = xp.broadcast_to(xp.asarray(True), (4, 1, 1))
+
+        actual = apply_where_jit(cond, self.f1, self.f2, (x, y))
+        expect = xp.where(cond, self.f1(x, y), self.f2(x, y))
+        xp_assert_equal(actual, expect)
+
+        actual = apply_where_jit(
+            cond,
+            lambda x, _: x,  # pyright: ignore[reportUnknownArgumentType]
+            lambda _, y: y,  # pyright: ignore[reportUnknownArgumentType]
+            (x, y),
+        )
+        expect = xp.where(cond, x, y)
+        xp_assert_equal(actual, expect)
+
+        # Shaped fill_value
+        actual = apply_where_jit(cond, self.f1, None, x, fill_value=y)
+        expect = xp.where(cond, self.f1(x), y)
+        xp_assert_equal(actual, expect)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_dtype_propagation(self, xp: ModuleType, library: Backend):
+        x = xp.asarray([1, 2], dtype=xp.int8)
+        y = xp.asarray([3, 4], dtype=xp.int16)
+        cond = x % 2 == 0
+
+        mxp = np if library is Backend.DASK else xp
+        actual = apply_where_jit(
+            cond,
+            self.f1,
+            lambda x, y: mxp.astype(x - y, xp.int64),  # pyright: ignore[reportUnknownArgumentType]
+            (x, y),
+        )
+        assert actual.dtype == xp.int64
+
+        actual = apply_where_jit(cond, self.f1, None, y, fill_value=5)
+        assert actual.dtype == xp.int16
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    @pytest.mark.parametrize("fill_value_raw", [3, [3, 4]])
+    @pytest.mark.parametrize(
+        ("fill_value_dtype", "expect_dtype"), [("int32", "int32"), ("int8", "int16")]
+    )
+    def test_dtype_propagation_fill_value(
+        self,
+        xp: ModuleType,
+        fill_value_raw: int | list[int],
+        fill_value_dtype: str,
+        expect_dtype: str,
+    ):
+        x = xp.asarray([1, 2], dtype=xp.int16)
+        cond = x % 2 == 0
+        fill_value = xp.asarray(fill_value_raw, dtype=getattr(xp, fill_value_dtype))
+
+        actual = apply_where_jit(cond, self.f1, None, x, fill_value=fill_value)
+        assert actual.dtype == getattr(xp, expect_dtype)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_dont_overwrite_fill_value(self, xp: ModuleType):
+        x = xp.asarray([1, 2])
+        fill_value = xp.asarray([100, 200])
+        actual = apply_where_jit(x % 2 == 0, self.f1, None, x, fill_value=fill_value)
+        xp_assert_equal(actual, xp.asarray([100, 12]))
+        xp_assert_equal(fill_value, xp.asarray([100, 200]))
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_dont_run_on_false(self, xp: ModuleType):
+        x = xp.asarray([1.0, 2.0, 0.0])
+        y = xp.asarray([0.0, 3.0, 4.0])
+        # On NumPy, division by zero will trigger warnings
+        actual = apply_where_jit(
+            x == 0,
+            lambda x, y: x / y,  # pyright: ignore[reportUnknownArgumentType]
+            lambda x, y: y / x,  # pyright: ignore[reportUnknownArgumentType]
+            (x, y),
+        )
+        xp_assert_equal(actual, xp.asarray([0.0, 1.5, 0.0]))
+
+    def test_bad_args(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4])
+        cond = x % 2 == 0
+        # Neither f2 nor fill_value
+        with pytest.raises(TypeError, match="Exactly one of"):
+            apply_where(cond, self.f1, x)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+        # Both f2 and fill_value
+        with pytest.raises(TypeError, match="Exactly one of"):
+            apply_where(cond, self.f1, self.f2, x, fill_value=0)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+        # Multiple args; forgot to wrap them in a tuple
+        with pytest.raises(TypeError, match="takes from 3 to 4 positional arguments"):
+            apply_where(cond, self.f1, self.f2, x, x)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+        with pytest.raises(TypeError, match="callable"):
+            apply_where(cond, self.f1, x, x, fill_value=0)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue]
+
+    @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_xp(self, xp: ModuleType):
+        x = xp.asarray([1, 2, 3, 4])
+        cond = x % 2 == 0
+        actual = apply_where_jit(cond, self.f1, self.f2, x, xp=xp)
+        expect = xp.where(cond, self.f1(x), self.f2(x))
+        xp_assert_equal(actual, expect)
+
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="read-only without .at")
+    def test_device(self, xp: ModuleType, device: Device):
+        x = xp.asarray([1, 2, 3, 4], device=device)
+        y = apply_where_jit(x % 2 == 0, self.f1, self.f2, x)
+        assert get_device(y) == device
+
+
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
 class TestAtLeastND:
     def test_0D(self, xp: ModuleType):
@@ -334,8 +509,6 @@ def test_xp(self, xp: ModuleType):
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
 class TestIsClose:
-    # FIXME use lazywhere to avoid warnings on inf
-    @pytest.mark.filterwarnings("ignore:invalid value encountered")
     @pytest.mark.parametrize("swap", [False, True])
     @pytest.mark.parametrize(
         ("a", "b"),
@@ -394,8 +567,6 @@ def test_broadcast(self, dtype: str, xp: ModuleType):
 
         xp_assert_equal(actual, expect)
 
-    # FIXME use lazywhere to avoid warnings on inf
-    @pytest.mark.filterwarnings("ignore:invalid value encountered")
     def test_some_inf(self, xp: ModuleType):
         a = xp.asarray([0.0, 1.0, xp.inf, xp.inf, xp.inf])
         b = xp.asarray([1e-9, 1.0, xp.inf, -xp.inf, 2.0])

tests/test_funcs.py

@@ -5,8 +5,9 @@
 
 from array_api_extra._lib import Backend
 from array_api_extra._lib._testing import xp_assert_equal
+from array_api_extra._lib._utils._compat import array_namespace
 from array_api_extra._lib._utils._compat import device as get_device
-from array_api_extra._lib._utils._helpers import asarrays, in1d
+from array_api_extra._lib._utils._helpers import asarrays, in1d, meta_namespace
 from array_api_extra._lib._utils._typing import Device
 from array_api_extra.testing import lazy_xp_function
 
@@ -15,6 +16,8 @@
 # FIXME calls xp.unique_values without size
 lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))
 
+np_compat = array_namespace(np.empty(0))
+
 
 class TestIn1D:
     @pytest.mark.xfail_xp_backend(
@@ -151,3 +154,21 @@ def test_asarrays_numpy_generics(dtype: type):
     xa, xb = asarrays(a, 0, xp=np)
     assert xa.dtype == dtype
     assert xb.dtype == dtype
+
+
+class TestMetaNamespace:
+    @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="namespace tests")
+    def test_basic(self, xp: ModuleType, library: Backend):
+        args = None, xp.asarray(0), 1
+        expect = np_compat if library is Backend.DASK else xp
+        assert meta_namespace(*args) is expect
+
+    def test_dask_metas(self, da: ModuleType):
+        cp = pytest.importorskip("cupy")
+        cp_compat = array_namespace(cp.empty(0))
+        args = None, da.from_array(cp.asarray(0)), 1
+        assert meta_namespace(*args) is cp_compat
+
+    def test_xp(self, xp: ModuleType):
+        args = None, xp.asarray(0), 1
+        assert meta_namespace(*args, xp=xp) in (xp, np_compat)