WIP

Author: crusaderky

src/array_api_extra/_lib/_lazy.py

@@ -9,6 +9,7 @@
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, cast, overload
 
+from ._utils import _compat
 from ._utils._compat import (
     array_namespace,
     is_array_api_obj,
@@ -295,7 +296,7 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         )
 
     else:
-        # Eager backends
+        # Eager backends, including non-jitted JAX
         wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
         out = wrapped(*args, **kwargs)
 
@@ -309,7 +310,7 @@ def _is_jax_jit_enabled(xp: ModuleType) -> bool:  # numpydoc ignore=PR01,RT01
     x = xp.asarray(False)
     try:
         return bool(x)
-    except jax.errors.TracerArrayConversionError:
+    except jax.errors.TracerBoolConversionError:
         return True
 
 
@@ -362,14 +363,16 @@ def _lazy_apply_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR0
     def wrapper(  # type: ignore[no-any-decorated,no-any-explicit]
         *args: Array, **kwargs: Any
     ) -> tuple[Array, ...]:  # numpydoc ignore=GL08
+        device = _compat.device(args[0]) if args else None
+
         if as_numpy:
             args = _as_numpy(args)
             kwargs = _as_numpy(kwargs)
         out = func(*args, **kwargs)
 
         if multi_output:
             assert isinstance(out, Sequence)
-            return tuple(xp.asarray(o) for o in out)
-        return (xp.asarray(out),)
+            return tuple(xp.asarray(o, device=device) for o in out)
+        return (xp.asarray(out, device=device),)
 
     return wrapper

tests/conftest.py

@@ -110,7 +110,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:
@@ -119,8 +123,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])  # Select the test with `pytest -k dask`
+def da(
+    request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> ModuleType:  # numpydoc ignore=PR01,RT01
+    """Fixture that returns dask.array (wrapped by array-api-compat)."""
+    xp = pytest.importorskip("dask.array")
+    xp = array_namespace(xp.empty(0))
+    patch_lazy_xp_functions(request, monkeypatch, xp=xp)
+    return xp
 
 
 @pytest.fixture

tests/test_lazy.py

@@ -10,51 +10,120 @@
 from array_api_extra._lib._utils._typing import Array
 from array_api_extra.testing import lazy_xp_function
 
-skip_as_numpy = [
-    pytest.mark.skip_xp_backend(Backend.CUPY, reason="device->host transfer"),
-    pytest.mark.skip_xp_backend(Backend.SPARSE, reason="densification"),
-]
+as_numpy = pytest.mark.parametrize(
+    "as_numpy",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=[
+                pytest.mark.skip_xp_backend(Backend.CUPY, reason="device->host copy"),
+                pytest.mark.skip_xp_backend(Backend.SPARSE, reason="densification"),
+            ],
+        ),
+    ],
+)
+
+
+@as_numpy
+def test_lazy_apply_simple(xp: ModuleType, as_numpy: bool):
+    pytest.skip("TODO")
+
+
+@as_numpy
+def test_lazy_apply_broadcast(xp: ModuleType, as_numpy: bool):
+    pytest.skip("TODO")
+
+
+@as_numpy
+def test_lazy_apply_multi_output(xp: ModuleType, as_numpy: bool):
+    pytest.skip("TODO")
+
+
+def test_lazy_apply_core_indices(da: ModuleType):
+    """Test that a func that performs reductions along axes does so
+    globally and not locally to each Dask chunk.
+    """
+    pytest.skip("TODO")
+
+
+def test_lazy_apply_dont_run_on_meta(da: ModuleType):
+    """Test that Dask won't try running func on the meta array,
+    as it may have minimum size requirements.
+    """
+    pytest.skip("TODO")
+
+
+def test_lazy_apply_none_shape(da: ModuleType):
+    pytest.skip("TODO")
+
+
+@as_numpy
+def test_lazy_apply_device(xp: ModuleType, as_numpy: bool):
+    pytest.skip("TODO")
+
+
+@as_numpy
+def test_lazy_apply_no_args(xp: ModuleType, as_numpy: bool):
+    pytest.skip("TODO")
 
 
 class NT(NamedTuple):
     a: Array
 
 
-def check_lazy_apply_kwargs(x: Array, expect: type, as_numpy: bool) -> Array:
-    def f(
+def check_lazy_apply_kwargs(x: Array, expect_cls: type, as_numpy: bool) -> Array:
+    def eager(
         x: Array,
         z: dict[str, list[Array] | tuple[Array, ...] | NT],
         msg: str,
         msgs: list[str],
+        scalar: int,
     ) -> Array:
-        assert isinstance(x, expect)
+        assert isinstance(x, expect_cls)
+        assert int(x) == 0  # JAX will crash if x isn't material
+        # Did we re-wrap the namedtuple correctly, or did it get
+        # accidentally changed to a basic tuple?
         assert isinstance(z["foo"], NT)
-        assert isinstance(z["foo"].a, expect)
-        assert isinstance(z["bar"][0], expect)
-        assert isinstance(z["baz"][0], expect)
-        assert msg == "Hello World"
-        assert msgs[0] == "Hello World"
-        return x + 1
+        assert isinstance(z["foo"].a, expect_cls)
+        assert isinstance(z["bar"][0], expect_cls)  # list
+        assert isinstance(z["baz"][0], expect_cls)  # tuple
+        assert msg == "Hello World"  # must be hidden from JAX
+        assert msgs[0] == "Hello World"  # must be hidden from JAX
+        assert isinstance(msg, str)
+        assert isinstance(msgs[0], str)
+        assert scalar == 1  # must be hidden from JAX
+        assert isinstance(scalar, int)
+        return x + 1  # type: ignore[operator]
 
     return lazy_apply(  # pyright: ignore[reportCallIssue]
-        f,
+        eager,
         x,
+        # These kwargs can and should be passed through jax.pure_callback
         z={"foo": NT(x), "bar": [x], "baz": (x,)},
+        # These can't
         msg="Hello World",
         msgs=["Hello World"],
+        # This will be automatically cast to jax.Array if we don't wrap it
+        scalar=1,
         shape=x.shape,
         dtype=x.dtype,
         as_numpy=as_numpy,
     )
 
-lazy_xp_function(check_lazy_apply_kwargs, static_argnames=("expect", "as_numpy"))
+
+lazy_xp_function(check_lazy_apply_kwargs, static_argnames=("expect_cls", "as_numpy"))
 
 
-@pytest.mark.parametrize("as_numpy", [False, pytest.param(True, marks=skip_as_numpy)])
+@as_numpy
 def test_lazy_apply_kwargs(xp: ModuleType, library: Backend, as_numpy: bool) -> None:
-    expect = np.ndarray if as_numpy or library is Backend.DASK else type(xp.asarray(0))
+    """When as_numpy=True, search and replace arrays in the (nested) keywords arguments
+    with numpy arrays, and leave the rest untouched."""
+    expect_cls = (
+        np.ndarray if as_numpy or library is Backend.DASK else type(xp.asarray(0))
+    )
     x = xp.asarray(0)
-    actual = check_lazy_apply_kwargs(x, expect, as_numpy)
+    actual = check_lazy_apply_kwargs(x, expect_cls, as_numpy)  # pyright: ignore[reportUnknownArgumentType]
     xp_assert_equal(actual, x + 1)
 
 
@@ -69,10 +138,12 @@ def eager(_: Array) -> Array:
 
     return lazy_apply(eager, x, shape=x.shape, dtype=x.dtype)
 
-lazy_xp_function(raises)
+
+# jax.pure_callback does not support raising
+# https://github.com/jax-ml/jax/issues/26102
+lazy_xp_function(raises, jax_jit=False)
 
 
-@pytest.mark.skip_xp_backend(Backend.JAX_JIT, reason="no exception support")
 def test_lazy_apply_raises(xp: ModuleType) -> None:
     x = xp.asarray(0)