WIP tests

Author: crusaderky

tests/conftest.py

@@ -148,19 +148,7 @@ def xp(
     patch_lazy_xp_functions(request, monkeypatch, xp=xp)
 
     if library.like(Backend.JAX):
-        import jax
-
-        # 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]
-
-        if library == Backend.JAX_GPU:
-            try:
-                device = jax.devices("cuda")[0]
-            except RuntimeError:
-                pytest.skip("no CUDA device available")
-        else:
-            device = jax.devices("cpu")[0]
-        jax.config.update("jax_default_device", device)
+        _setup_jax(library)
 
     elif library == Backend.TORCH_GPU:
         import torch.cuda
@@ -175,6 +163,22 @@ def xp(
     yield xp
 
 
+def _setup_jax(library: Backend) -> None:
+    import jax
+
+    # 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]
+
+    if library == Backend.JAX_GPU:
+        try:
+            device = jax.devices("cuda")[0]
+        except RuntimeError:
+            pytest.skip("no CUDA device available")
+    else:
+        device = jax.devices("cpu")[0]
+    jax.config.update("jax_default_device", device)
+
+
 @pytest.fixture(params=[Backend.DASK])  # Can select the test with `pytest -k dask`
 def da(
     request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
@@ -186,6 +190,17 @@ def da(
     return xp
 
 
+@pytest.fixture(params=[Backend.JAX, Backend.JAX_GPU])
+def jnp(
+    request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> ModuleType:  # numpydoc ignore=PR01,RT01
+    """Variant of the `xp` fixture that only yields jax.numpy."""
+    xp = pytest.importorskip("jax.numpy")
+    _setup_jax(request.param)
+    patch_lazy_xp_functions(request, monkeypatch, xp=xp)
+    return xp
+
+
 @pytest.fixture
 def device(
     library: Backend, xp: ModuleType

tests/test_helpers.py

@@ -1,5 +1,5 @@
 from types import ModuleType
-from typing import cast
+from typing import Generic, TypeVar, cast
 
 import numpy as np
 import pytest
@@ -13,6 +13,7 @@
     capabilities,
     eager_shape,
     in1d,
+    jax_autojit,
     meta_namespace,
     ndindex,
 )
@@ -23,6 +24,8 @@
 
 # mypy: disable-error-code=no-untyped-usage
 
+T = TypeVar("T")
+
 # FIXME calls xp.unique_values without size
 lazy_xp_function(in1d, jax_jit=False)
 
@@ -204,3 +207,90 @@ def test_capabilities(xp: ModuleType):
     if xp.__array_api_version__ >= "2024.12":
         expect.add("max dimensions")
     assert capabilities(xp).keys() == expect
+
+
+class Wrapper(Generic[T]):
+    """Trivial opaque wrapper. Must be pickleable."""
+
+    x: T
+    # __slots__ make this object serializable with __reduce_ex__(5),
+    # but not with __reduce__
+    __slots__: tuple[str, ...] = ("x",)
+
+    def __init__(self, x: T):
+        self.x = x
+
+
+class TestJAXAutoJIT:
+    def test_basic(self, jnp: ModuleType):
+        @jax_autojit
+        def f(x: Array, k: object = False) -> Array:
+            return x + 1 if k else x - 1
+
+        # Basic recognition of static_argnames
+        xp_assert_equal(f(jnp.asarray([1, 2])), jnp.asarray([0, 1]))
+        xp_assert_equal(f(jnp.asarray([1, 2]), False), jnp.asarray([0, 1]))
+        xp_assert_equal(f(jnp.asarray([1, 2]), True), jnp.asarray([2, 3]))
+        xp_assert_equal(f(jnp.asarray([1, 2]), 1), jnp.asarray([2, 3]))
+
+        # static argument is not an ArrayLike
+        xp_assert_equal(f(jnp.asarray([1, 2]), "foo"), jnp.asarray([2, 3]))
+
+        # static argument is not hashable, but serializable
+        xp_assert_equal(f(jnp.asarray([1, 2]), ["foo"]), jnp.asarray([2, 3]))
+
+    def test_wrapper(self, jnp: ModuleType):
+        @jax_autojit
+        def f(w: Wrapper[Array]) -> Wrapper[Array]:
+            return Wrapper(w.x + 1)
+
+        inp = Wrapper(jnp.asarray([1, 2]))
+        out = f(inp).x
+        xp_assert_equal(out, jnp.asarray([2, 3]))
+
+    def test_static_hashable(self, jnp: ModuleType):
+        """Static argument/return value is hashable, but not serializable"""
+
+        class C:
+            def __reduce__(self) -> object:  # type: ignore[explicit-override,override]  # pyright: ignore[reportIncompatibleMethodOverride,reportImplicitOverride]
+                raise Exception()
+
+        @jax_autojit
+        def f(x: object) -> object:
+            return x
+
+        inp = C()
+        out = f(inp)
+        assert out is inp
+
+        # Serializable opaque input contains non-serializable object plus array
+        inp = Wrapper((C(), jnp.asarray([1, 2])))
+        out = f(inp)
+        assert isinstance(out, Wrapper)
+        assert out.x[0] is inp.x[0]
+        assert out.x[1] is not inp.x[1]
+        xp_assert_equal(out.x[1], inp.x[1])  # pyright: ignore[reportUnknownArgumentType]
+
+    def test_arraylikes_are_static(self):
+        pytest.importorskip("jax")
+
+        @jax_autojit
+        def f(x: list[int]) -> list[int]:
+            assert isinstance(x, list)
+            assert x == [1, 2]
+            return [3, 4]
+
+        out = f([1, 2])
+        assert isinstance(out, list)
+        assert out == [3, 4]
+
+    def test_repeated_objects(self, jnp: ModuleType):
+        @jax_autojit
+        def f(x: Array, y: Array) -> tuple[Array, Array]:
+            z = x + y
+            return z, z
+
+        inp = jnp.asarray([1, 2])
+        out = f(inp, inp)
+        assert out[0] is out[1]
+        xp_assert_equal(out[0], jnp.asarray([2, 4]))