ENH: testing.lazy_xp_function: torch.compile support

Author: lucascolley

pyproject.toml

@@ -146,7 +146,11 @@ errors = { unannotated-return = false }
 
 [tool.pytest.ini_options]
 addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
-filterwarnings = ["error"]
+filterwarnings = [
+  "error",
+  "ignore:.*torch.jit.script_method.*",
+  "ignore:.*accumulated_recompile_limit reached.*",
+]
 log_cli_level = "INFO"
 markers = [
   "skip_xp_backend(library, /, *, reason=None): Skip test for a specific backend",

src/array_api_extra/_lib/_helpers.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import enum
 import functools
 import io
 import math
@@ -28,14 +29,15 @@
 
 
 __all__ = [
+    "JitLibrary",
     "asarrays",
+    "autojit",
     "capabilities",
     "deprecated",
     "eager_shape",
     "in1d",
     "is_jax_jit_enabled",
     "is_python_scalar",
-    "jax_autojit",
     "meta_namespace",
     "normalize_pad_width",
     "pickle_flatten",
@@ -487,20 +489,20 @@ def persistent_load(self, pid: Literal[0, 1]) -> object:  # numpydoc ignore=GL08
 
 class _AutoJITWrapper(Generic[T]):  # numpydoc ignore=PR01
     """
-    Helper of :func:`jax_autojit`.
+    Helper of :func:`autojit`.
 
     Wrap arbitrary inputs and outputs of the jitted function and
     convert them to/from PyTrees.
     """
 
     _obj: Any
     _is_iter: bool
-    _registered: ClassVar[bool] = False
+    _registered: ClassVar[set[JitLibrary]] = set()
     __slots__: tuple[str, ...] = ("_is_iter", "_obj")
 
-    def __init__(self, obj: T) -> None:  # numpydoc ignore=GL08
-        self._register()
-        if isinstance(obj, Iterator):
+    def __init__(self, obj: T, jit_library: JitLibrary) -> None:  # numpydoc ignore=GL08
+        self._register(jit_library)
+        if jit_library is JitLibrary.jax and isinstance(obj, Iterator):
             self._obj = list(obj)
             self._is_iter = True
         else:
@@ -513,24 +515,44 @@ def obj(self) -> T:  # numpydoc ignore=RT01
         return iter(self._obj) if self._is_iter else self._obj
 
     @classmethod
-    def _register(cls) -> None:  # numpydoc ignore=SS06
+    def _register(cls, jit_library: JitLibrary) -> None:  # numpydoc ignore=SS06,PR01
         """
         Register upon first use instead of at import time, to avoid
         globally importing JAX.
         """
-        if not cls._registered:
+        if jit_library in cls._registered:
+            return
+
+        if jit_library is JitLibrary.jax:
             import jax
 
             jax.tree_util.register_pytree_node(
                 cls,
                 lambda instance: pickle_flatten(instance, jax.Array),  # pyright: ignore[reportUnknownArgumentType]
                 lambda aux_data, children: pickle_unflatten(children, aux_data),  # pyright: ignore[reportUnknownArgumentType]
             )
-            cls._registered = True
+        elif jit_library is JitLibrary.torch:
+            import torch
+
+            torch.utils._pytree.register_pytree_node(
+                cls,
+                lambda instance: pickle_flatten(instance, torch.Tensor),  # pyright: ignore[reportUnknownArgumentType]
+                pickle_unflatten,
+            )
+        cls._registered.add(jit_library)
 
 
-def jax_autojit(
-    func: Callable[P, T],
+class JitLibrary(enum.Enum):
+    """
+    Enum for JIT libraries compatible with `autojit`.
+    """
+
+    jax = enum.auto()
+    torch = enum.auto()
+
+
+def autojit(
+    func: Callable[P, T], jit_library: JitLibrary
 ) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01,SS03
     """
     Wrap `func` with ``jax.jit``, with the following differences:
@@ -573,19 +595,26 @@ def f(x: Array, y: float, plus: bool) -> Array:
     ``j1``, but on the flip side it means that it will be re-traced for every different
     value of ``y``, which likely makes it not fit for purpose in production.
     """
-    import jax
+    if jit_library is JitLibrary.jax:
+        import jax
+
+        jit_decorator = jax.jit
+    elif jit_library is JitLibrary.torch:
+        import torch
+
+        jit_decorator = functools.partial(torch.compile, fullgraph=True)
 
-    @jax.jit  # type: ignore[untyped-decorator]  # pyright: ignore[reportUntypedFunctionDecorator]
+    @jit_decorator  # type: ignore[untyped-decorator]  # pyright: ignore[reportUntypedFunctionDecorator]
     def inner(  # numpydoc ignore=GL08
         wargs: _AutoJITWrapper[Any],
     ) -> _AutoJITWrapper[T]:
         args, kwargs = wargs.obj
         res = func(*args, **kwargs)  # pyright: ignore[reportCallIssue]
-        return _AutoJITWrapper(res)
+        return _AutoJITWrapper(res, jit_library)
 
     @functools.wraps(func)
     def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
-        wargs = _AutoJITWrapper((args, kwargs))
+        wargs = _AutoJITWrapper((args, kwargs), jit_library)
         return inner(wargs).obj
 
     return outer

src/array_api_extra/testing/_testing.py

@@ -37,17 +37,6 @@
     "patch_lazy_xp_functions",
 ]
 
-
-__all__ = [
-    "assert_close",
-    "assert_close_nulp",
-    "assert_equal",
-    "assert_less",
-    "lazy_xp_function",
-    "patch_lazy_xp_functions",
-]
-
-
 P = ParamSpec("P")
 T = TypeVar("T")
 
@@ -83,6 +72,7 @@ def lazy_xp_function(
     *,
     allow_dask_compute: bool | int = False,
     jax_jit: bool = True,
+    torch_compile: bool = True,
     static_argnums: _Deprecated = DEPRECATED,
     static_argnames: _Deprecated = DEPRECATED,
 ) -> None:  # numpydoc ignore=GL07
@@ -146,6 +136,8 @@ def lazy_xp_function(
             ...     return user_consumes(z)
 
         Default: True.
+    torch_compile : bool, optional
+        TODO: proper docs.
     static_argnums : Deprecated
         Deprecated; ignored.
     static_argnames : Deprecated
@@ -238,6 +230,7 @@ def test_myfunc(xp):
     tags: dict[str, bool | int | type] = {
         "allow_dask_compute": allow_dask_compute,
         "jax_jit": jax_jit,
+        "torch_compile": torch_compile,
     }
 
     if isinstance(func, tuple):
@@ -444,7 +437,19 @@ def iter_tagged() -> Iterator[
     elif _compat.is_jax_namespace(xp):
         for target, name, attr, func, tags in iter_tagged():
             if tags["jax_jit"]:
-                wrapped = _helpers.jax_autojit(func)
+                wrapped = _helpers.autojit(func, _helpers.JitLibrary.jax)
+                # If we're dealing with a staticmethod or classmethod, make
+                # sure things stay that way.
+                if isinstance(attr, staticmethod):
+                    wrapped = staticmethod(wrapped)
+                elif isinstance(attr, classmethod):
+                    wrapped = classmethod(wrapped)
+                temp_setattr(target, name, wrapped)
+
+    elif _compat.is_torch_namespace(xp):
+        for target, name, attr, func, tags in iter_tagged():
+            if tags["torch_compile"]:
+                wrapped = _helpers.autojit(func, _helpers.JitLibrary.torch)
                 # If we're dealing with a staticmethod or classmethod, make
                 # sure things stay that way.
                 if isinstance(attr, staticmethod):

tests/main/test_helpers.py

@@ -1,5 +1,7 @@
-from collections.abc import Iterator
-from typing import Generic, TypeVar, cast
+import functools
+from collections.abc import Callable, Iterator
+from types import ModuleType
+from typing import Generic, ParamSpec, Protocol, TypeVar, cast
 
 import numpy as np
 import pytest
@@ -8,11 +10,12 @@
 from array_api_extra._lib._compat import array_namespace
 from array_api_extra._lib._compat import device as get_device
 from array_api_extra._lib._helpers import (
+    JitLibrary,
     asarrays,
+    autojit,
     capabilities,
     eager_shape,
     in1d,
-    jax_autojit,
     meta_namespace,
     ndindex,
     pickle_flatten,
@@ -23,6 +26,7 @@
 
 from .conftest import np_compat
 
+P = ParamSpec("P")
 T = TypeVar("T")
 
 # FIXME calls xp.unique_values without size
@@ -345,41 +349,48 @@ def test_recursion(self):
         assert obj2[1] is obj2
 
 
-class TestJAXAutoJIT:
-    def test_basic(self, jnp: ArrayNamespace):
-        @jax_autojit
+class AutoJitFunc(Protocol):
+    def __call__(
+        self,
+        func: Callable[P, T],
+    ) -> Callable[P, T]: ...
+
+
+class CheckAutoJIT:
+    def test_basic(self, autojit_func: AutoJitFunc, xp: ArrayNamespace):
+        @autojit_func
         def f(x: Array, k: object = False) -> Array:
             return x + 1 if k else x - 1
 
         # Basic recognition of static_argnames
-        assert_equal(f(jnp.asarray([1, 2])), jnp.asarray([0, 1]))
-        assert_equal(f(jnp.asarray([1, 2]), False), jnp.asarray([0, 1]))
-        assert_equal(f(jnp.asarray([1, 2]), True), jnp.asarray([2, 3]))
-        assert_equal(f(jnp.asarray([1, 2]), 1), jnp.asarray([2, 3]))
+        assert_equal(f(xp.asarray([1, 2])), xp.asarray([0, 1]))
+        assert_equal(f(xp.asarray([1, 2]), False), xp.asarray([0, 1]))
+        assert_equal(f(xp.asarray([1, 2]), True), xp.asarray([2, 3]))
+        assert_equal(f(xp.asarray([1, 2]), 1), xp.asarray([2, 3]))
 
         # static argument is not an ArrayLike
-        assert_equal(f(jnp.asarray([1, 2]), "foo"), jnp.asarray([2, 3]))
+        assert_equal(f(xp.asarray([1, 2]), "foo"), xp.asarray([2, 3]))
 
         # static argument is not hashable, but serializable
-        assert_equal(f(jnp.asarray([1, 2]), ["foo"]), jnp.asarray([2, 3]))
+        assert_equal(f(xp.asarray([1, 2]), ["foo"]), xp.asarray([2, 3]))
 
-    def test_wrapper(self, jnp: ArrayNamespace):
-        @jax_autojit
+    def test_wrapper(self, autojit_func: AutoJitFunc, xp: ArrayNamespace):
+        @autojit_func
         def f(w: Wrapper[Array]) -> Wrapper[Array]:
             return Wrapper(w.x + 1)
 
-        inp = Wrapper(jnp.asarray([1, 2]))
+        inp = Wrapper(xp.asarray([1, 2]))
         out = f(inp).x
-        assert_equal(out, jnp.asarray([2, 3]))
+        assert_equal(out, xp.asarray([2, 3]))
 
-    def test_static_hashable(self, jnp: ArrayNamespace):
+    def test_static_hashable(self, autojit_func: AutoJitFunc, xp: ArrayNamespace):
         """Static argument/return value is hashable, but not serializable"""
 
         class C:
             def __reduce__(self) -> object:  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride,reportImplicitOverride]
                 raise Exception()  # noqa: TRY002
 
-        @jax_autojit
+        @autojit_func
         def f(x: object) -> object:
             return x
 
@@ -388,17 +399,20 @@ def f(x: object) -> object:
         assert out is inp
 
         # Serializable opaque input contains non-serializable object plus array
-        winp = Wrapper((C(), jnp.asarray([1, 2])))
+        winp = Wrapper((C(), xp.asarray([1, 2])))
         out = f(winp)
         assert isinstance(out, Wrapper)
         assert out.x[0] is winp.x[0]
         assert out.x[1] is not winp.x[1]
         assert_equal(out.x[1], winp.x[1])
 
-    def test_arraylikes_are_static(self):
+    def test_arraylikes_are_static(
+        self,
+        autojit_func: AutoJitFunc,
+    ):
         pytest.importorskip("jax")
 
-        @jax_autojit
+        @autojit_func
         def f(x: list[int]) -> list[int]:
             assert isinstance(x, list)
             assert x == [1, 2]
@@ -408,15 +422,35 @@ def f(x: list[int]) -> list[int]:
         assert isinstance(out, list)
         assert out == [3, 4]
 
-    def test_iterators(self, jnp: ArrayNamespace):
-        @jax_autojit
+    def test_iterators(self, autojit_func: AutoJitFunc, xp: ArrayNamespace):
+        @autojit_func
         def f(x: Array) -> Iterator[Array]:
             return (x + i for i in range(2))
 
-        inp = jnp.asarray([1, 2])
+        inp = xp.asarray([1, 2])
         out = f(inp)
         assert isinstance(out, Iterator)
-        assert_equal(next(out), jnp.asarray([1, 2]))
-        assert_equal(next(out), jnp.asarray([2, 3]))
+        assert_equal(next(out), xp.asarray([1, 2]))
+        assert_equal(next(out), xp.asarray([2, 3]))
         with pytest.raises(StopIteration):
             _ = next(out)
+
+
+class TestJAXAutoJit(CheckAutoJIT):
+    @pytest.fixture
+    def xp(self, jnp: ModuleType) -> ModuleType:
+        return jnp
+
+    @pytest.fixture
+    def autojit_func(self) -> AutoJitFunc:
+        return functools.partial(autojit, jit_library=JitLibrary.jax)
+
+
+class TestTorchAutoJit(CheckAutoJIT):
+    @pytest.fixture
+    def xp(self, torch: ModuleType) -> ModuleType:
+        return torch
+
+    @pytest.fixture
+    def autojit_func(self) -> AutoJitFunc:
+        return functools.partial(autojit, jit_library=JitLibrary.torch)