ENH: testing.lazy_xp_function: torch.compile support

.github/workflows/ci.yml

@@ -68,7 +68,7 @@ jobs:
           # Non-default pixi task (default: tests-ci)
           - {
               environment: tests-nogil,
-              task: "tests --parallel-threads=4",
+              task: tests "--parallel-threads=4 tests/main",
               no-coverage: true,
               x64_runner: ubuntu-latest,
             }
@@ -93,7 +93,7 @@ jobs:
           environments: ${{ matrix.environment }}
 
       - name: Test
-        run: pixi run -e "${TASK_ENV}" ${TASK}
+        run: pixi run -e "${TASK_ENV}" "${TASK}"
         env:
           TASK_ENV: ${{ matrix.environment }}
           TASK: ${{ matrix.task || 'tests-ci' }}

pixi.toml

@@ -224,12 +224,14 @@ scipy = ">=1.15.2"
 [feature.tests.tasks]
 tests = {
   description = "Run tests",
-  cmd = "pytest -v tests/main",
+  args = [{ arg = "pytest_args", default = "tests/main" }],
+  cmd = "pytest -v {{ pytest_args }}",
   default-environment = "tests",
 }
 tests-cov = {
   description = "Run tests with coverage",
-  cmd = "pytest -v -ra --cov --cov-report=xml --cov-report=term --durations=20 tests/main",
+  args = [{ arg = "pytest_args", default = "tests/main" }],
+  cmd = "pytest -v -ra --cov --cov-report=xml --cov-report=term --durations=20 {{ pytest_args }}",
   default-environment = "tests",
 }
 
@@ -361,6 +363,7 @@ numpy = "=1.24.1"
 # Note: JAX and PyTorch will install CPU variants.
 [feature.backends.dependencies]
 pytorch = ">=2.12.0"
+cxx-compiler = ">=1.11.0,<2"  # for torch.compile
 dask-core = ">=2026.7.1"  # No distributed, tornado, etc.
 sparse = ">=0.19.0"
 

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:.*dynamo.*functools.lru_cache.*",
+]
 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,27 @@ 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)
+        jit_decorator = functools.partial(torch.compile, fullgraph=False)
 
-    @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):