jax_autojit

Author: crusaderky

src/array_api_extra/testing.py

@@ -7,7 +7,9 @@
 from __future__ import annotations
 
 import contextlib
-from collections.abc import Callable, Iterable, Iterator, Sequence
+import enum
+import warnings
+from collections.abc import Callable, Hashable, Iterator, Sequence
 from functools import wraps
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
@@ -37,13 +39,22 @@ def override(func: object) -> object:
 _ufuncs_tags: dict[object, dict[str, Any]] = {}  # type: ignore[explicit-any]
 
 
+class Deprecated(enum.Enum):
+    """Unique type for deprecated parameters."""
+
+    DEPRECATED = 1
+
+
+DEPRECATED = Deprecated.DEPRECATED
+
+
 def lazy_xp_function(  # type: ignore[explicit-any]
     func: Callable[..., Any],
     *,
     allow_dask_compute: int = 0,
     jax_jit: bool = True,
-    static_argnums: int | Sequence[int] | None = None,
-    static_argnames: str | Iterable[str] | None = None,
+    static_argnums: Deprecated = DEPRECATED,
+    static_argnames: Deprecated = DEPRECATED,
 ) -> None:  # numpydoc ignore=GL07
     """
     Tag a function to be tested on lazy backends.
@@ -79,16 +90,15 @@ def lazy_xp_function(  # type: ignore[explicit-any]
         Default: 0, meaning that `func` must be fully lazy and never materialize the
         graph.
     jax_jit : bool, optional
-        Set to True to replace `func` with ``jax.jit(func)`` after calling the
-        :func:`patch_lazy_xp_functions` test helper with ``xp=jax.numpy``. Set to False
-        if `func` is only compatible with eager (non-jitted) JAX. Default: True.
-    static_argnums : int | Sequence[int], optional
-        Passed to jax.jit. Positional arguments to treat as static (compile-time
-        constant). Default: infer from `static_argnames` using
-        `inspect.signature(func)`.
-    static_argnames : str | Iterable[str], optional
-        Passed to jax.jit. Named arguments to treat as static (compile-time constant).
-        Default: infer from `static_argnums` using `inspect.signature(func)`.
+        Set to True to replace `func` with a variant of ``jax.jit(func)``
+        (read notes below) after calling the :func:`patch_lazy_xp_functions`
+        test helper with ``xp=jax.numpy``.
+        Set to False if `func` is only compatible with eager (non-jitted) JAX.
+        Default: True.
+    static_argnums :
+        Deprecated; ignored
+    static_argnames :
+        Deprecated; ignored
 
     See Also
     --------
@@ -165,12 +175,20 @@ def test_myfunc(xp):
           b = mymodule.myfunc(a)  # This is wrapped when xp=jax.numpy or xp=dask.array
           c = naked.myfunc(a)  # This is not
     """
+    if static_argnums is not DEPRECATED or static_argnames is not DEPRECATED:
+        warnings.warn(
+            (
+                "The `static_argnums` and `static_argnames` parameters are deprecated "
+                "and ignored. They will be removed in a future version."
+            ),
+            DeprecationWarning,
+            stacklevel=2,
+        )
     tags = {
         "allow_dask_compute": allow_dask_compute,
         "jax_jit": jax_jit,
-        "static_argnums": static_argnums,
-        "static_argnames": static_argnames,
     }
+
     try:
         func._lazy_xp_function = tags  # type: ignore[attr-defined]  # pylint: disable=protected-access  # pyright: ignore[reportFunctionMemberAccess]
     except AttributeError:  # @cython.vectorize
@@ -240,15 +258,9 @@ def iter_tagged() -> (  # type: ignore[explicit-any]
             monkeypatch.setattr(mod, name, wrapped)
 
     elif is_jax_namespace(xp):
-        import jax
-
         for mod, name, func, tags in iter_tagged():
             if tags["jax_jit"]:
-                wrapped = _jax_wrap(
-                    func,
-                    static_argnums=tags["static_argnums"],
-                    static_argnames=tags["static_argnames"],
-                )
+                wrapped = _jax_autojit(func)
                 monkeypatch.setattr(mod, name, wrapped)
 
 
@@ -319,38 +331,50 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         # not work on scheduler='distributed', as it would not block.
         pik, arrays = pickle_without(out, da.Array)
         arrays = dask.persist(arrays, scheduler="threads")[0]  # type: ignore[attr-defined,no-untyped-call,func-returns-value,index]  # pyright: ignore[reportPrivateImportUsage]
-        unpickle_without(pik, arrays)
+        return unpickle_without(pik, arrays)
 
     return wrapper
 
 
-def _jax_wrap(
-    func: Callable[P, T],
-    static_argnums: int | Sequence[int] | None,
-    static_argnames: str | Iterable[str] | None,
-) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01
+def _jax_autojit(func: Callable[P, T]) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01
     """
-    Wrap `func` inside ``jax.jit``.
-
-    Accepts non-array return values.
+    Wrap `func` with ``jax.jit``, with the following differences:
+
+    - Array-like arguments and return values are not automatically converted to
+      ``jax.Array`` objects.
+    - All non-array arguments are automatically treated as static.
+    - Unlike ``jax.jit``, non-array arguments and return values are not limited to
+      tuple/list/dict, but can be any object serializable with ``pickle``.
+    - Automatically descend into non-array arguments and find ``jax.Array`` objects
+      inside them.
+    - Automatically descend into non-array return values and find ``jax.Array`` objects
+      inside them, then rebuild them downstream of exiting the JIT, swapping the JAX
+      tracer objects with concrete arrays.
     """
     import jax
-    import jax.numpy as jnp
 
-    def inner(*args: P.args, **kwargs: P.kwargs) -> tuple[jax.Array, ...]:  # numpydoc ignore=GL08
-        out = func(*args, **kwargs)
-        pik, arrays = pickle_without(out, jax.Array)
-        return jnp.frombuffer(pik, dtype=jnp.uint8), *arrays
+    # pickled return values of `func`, minus the JAX arrays
+    res_piks = {}
 
-    jitted = jax.jit(
-        inner, static_argnums=static_argnums, static_argnames=static_argnames
-    )
-    cpu = jax.devices("cpu")[0]
+    def jit_cache_key(args_pik: bytes, *arg_arrays: jax.Array) -> Hashable:  # type: ignore[no-any-unimported]
+        return (args_pik, *((arr.shape, arr.dtype) for arr in arg_arrays))
+
+    def inner(args_pik: bytes, *arg_arrays: jax.Array) -> list[jax.Array]:  # type: ignore[no-any-unimported]  # numpydoc ignore=GL08
+        args, kwargs = unpickle_without(args_pik, arg_arrays)
+        res = func(*args, **kwargs)
+        res_pik, res_arrays = pickle_without(res, jax.Array)
+        key = jit_cache_key(args_pik, *arg_arrays)
+        prev = res_piks.setdefault(key, res_pik)
+        assert prev == res_pik, "cache key collision"
+        return res_arrays
+
+    jitted = jax.jit(inner, static_argnums=0)
 
     @wraps(func)
     def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
-        arrays = jitted(*args, **kwargs)
-        pik = bytes(arrays[0].to_device(cpu))
-        return unpickle_without(pik, arrays[1:])
+        args_pik, arg_arrays = pickle_without((args, kwargs), jax.Array)
+        res_arrays = jitted(args_pik, *arg_arrays)
+        res_pik = res_piks[jit_cache_key(args_pik, *arg_arrays)]
+        return unpickle_without(res_pik, res_arrays)
 
     return outer

tests/test_at.py

@@ -1,5 +1,4 @@
 import math
-import pickle
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from types import ModuleType
@@ -41,28 +40,11 @@ def at_op(
     just a workaround for when one wants to apply jax.jit to `at()` directly,
     which is not a common use case.
     """
-    if isinstance(idx, (slice | tuple)):
-        return _at_op(x, None, pickle.dumps(idx), op, y, copy=copy, xp=xp)
-    return _at_op(x, idx, None, op, y, copy=copy, xp=xp)
-
-
-def _at_op(
-    x: Array,
-    idx: SetIndex | None,
-    idx_pickle: bytes | None,
-    op: _AtOp,
-    y: Array | object,
-    copy: bool | None,
-    xp: ModuleType | None = None,
-) -> Array:
-    """jitted helper of at_op"""
-    if idx_pickle:
-        idx = pickle.loads(idx_pickle)
-    meth = cast(Callable[..., Array], getattr(at(x, cast(SetIndex, idx)), op.value))  # type: ignore[explicit-any]
+    meth = cast(Callable[..., Array], getattr(at(x, idx), op.value))  # type: ignore[explicit-any]
     return meth(y, copy=copy, xp=xp)
 
 
-lazy_xp_function(_at_op, static_argnames=("op", "idx_pickle", "copy", "xp"))
+lazy_xp_function(at_op)
 
 
 @contextmanager

tests/test_funcs.py

@@ -37,17 +37,17 @@
 # some xp backends are untyped
 # mypy: disable-error-code=no-untyped-def
 
-lazy_xp_function(apply_where, static_argnums=(2, 3), static_argnames="xp")
-lazy_xp_function(atleast_nd, static_argnames=("ndim", "xp"))
-lazy_xp_function(cov, static_argnames="xp")
-lazy_xp_function(create_diagonal, static_argnames=("offset", "xp"))
-lazy_xp_function(expand_dims, static_argnames=("axis", "xp"))
-lazy_xp_function(kron, static_argnames="xp")
-lazy_xp_function(nunique, static_argnames="xp")
-lazy_xp_function(pad, static_argnames=("pad_width", "mode", "constant_values", "xp"))
+lazy_xp_function(apply_where)
+lazy_xp_function(atleast_nd)
+lazy_xp_function(cov)
+lazy_xp_function(create_diagonal)
+lazy_xp_function(expand_dims)
+lazy_xp_function(kron)
+lazy_xp_function(nunique)
+lazy_xp_function(pad)
 # FIXME calls in1d which calls xp.unique_values without size
-lazy_xp_function(setdiff1d, jax_jit=False, static_argnames=("assume_unique", "xp"))
-lazy_xp_function(sinc, static_argnames="xp")
+lazy_xp_function(setdiff1d, jax_jit=False)
+lazy_xp_function(sinc)
 
 
 class TestApplyWhere:

tests/test_helpers.py

@@ -24,7 +24,7 @@
 # mypy: disable-error-code=no-untyped-usage
 
 # FIXME calls xp.unique_values without size
-lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))
+lazy_xp_function(in1d, jax_jit=False)
 
 
 @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no unique_inverse")

tests/test_lazy.py

@@ -15,9 +15,7 @@
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
 
-lazy_xp_function(
-    lazy_apply, static_argnames=("func", "shape", "dtype", "as_numpy", "xp")
-)
+lazy_xp_function(lazy_apply)
 
 as_numpy = pytest.mark.parametrize(
     "as_numpy",
@@ -386,7 +384,7 @@ def eager(
     )
 
 
-lazy_xp_function(check_lazy_apply_kwargs, static_argnames=("expect_cls", "as_numpy"))
+lazy_xp_function(check_lazy_apply_kwargs)
 
 
 @as_numpy

tests/test_testing.py

@@ -179,32 +179,24 @@ def static_params(x: Array, n: int, flag: bool = False) -> Array:
     return x * 3.0
 
 
-def static_params1(x: Array, n: int, flag: bool = False) -> Array:
-    return static_params(x, n, flag)
+lazy_xp_function(static_params)
 
 
-def static_params2(x: Array, n: int, flag: bool = False) -> Array:
-    return static_params(x, n, flag)
-
-
-def static_params3(x: Array, n: int, flag: bool = False) -> Array:
-    return static_params(x, n, flag)
-
-
-lazy_xp_function(static_params1, static_argnums=(1, 2))
-lazy_xp_function(static_params2, static_argnames=("n", "flag"))
-lazy_xp_function(static_params3, static_argnums=1, static_argnames="flag")
-
-
-@pytest.mark.parametrize("func", [static_params1, static_params2, static_params3])
-def test_lazy_xp_function_static_params(xp: ModuleType, func: Callable[..., Array]):  # type: ignore[explicit-any]
+def test_lazy_xp_function_static_params(xp: ModuleType):
     x = xp.asarray([1.0, 2.0])
-    xp_assert_equal(func(x, 1), xp.asarray([3.0, 6.0]))
-    xp_assert_equal(func(x, 1, True), xp.asarray([2.0, 4.0]))
-    xp_assert_equal(func(x, 1, False), xp.asarray([3.0, 6.0]))
-    xp_assert_equal(func(x, 0, False), xp.asarray([3.0, 6.0]))
-    xp_assert_equal(func(x, 1, flag=True), xp.asarray([2.0, 4.0]))
-    xp_assert_equal(func(x, n=1, flag=True), xp.asarray([2.0, 4.0]))
+    xp_assert_equal(static_params(x, 1), xp.asarray([3.0, 6.0]))
+    xp_assert_equal(static_params(x, 1, True), xp.asarray([2.0, 4.0]))
+    xp_assert_equal(static_params(x, 1, False), xp.asarray([3.0, 6.0]))
+    xp_assert_equal(static_params(x, 0, False), xp.asarray([3.0, 6.0]))
+    xp_assert_equal(static_params(x, 1, flag=True), xp.asarray([2.0, 4.0]))
+    xp_assert_equal(static_params(x, n=1, flag=True), xp.asarray([2.0, 4.0]))
+
+
+def test_lazy_xp_function_deprecated_static_argnames():
+    with pytest.warns(DeprecationWarning, match="static_argnames"):
+        lazy_xp_function(static_params, static_argnames=["flag"])  # type: ignore[arg-type]
+    with pytest.warns(DeprecationWarning, match="static_argnums"):
+        lazy_xp_function(static_params, static_argnums=[1])  # type: ignore[arg-type]
 
 
 try: