jax_autojit

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-import copyreg
 import io
 import math
 import pickle
-from collections.abc import Callable, Generator, Iterable, Iterator
-from contextvars import ContextVar
-from types import ModuleType
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from collections.abc import Callable, Generator, Hashable, Iterable
+from functools import wraps
+from types import ModuleType, NoneType
+from typing import TYPE_CHECKING, Any, Literal, ParamSpec, TypeVar, cast
 
 from . import _compat
 from ._compat import (
@@ -23,9 +22,15 @@
 from ._typing import Array
 
 if TYPE_CHECKING:  # pragma: no cover
-    # TODO import from typing (requires Python >=3.13)
-    from typing_extensions import TypeIs
+    # TODO import from typing (requires Python >=3.12 and >=3.13)
+    from typing_extensions import TypeIs, override
+else:
 
+    def override(func):
+        return func
+
+
+P = ParamSpec("P")
 T = TypeVar("T")
 
 
@@ -35,6 +40,7 @@
     "eager_shape",
     "in1d",
     "is_python_scalar",
+    "jax_autojit",
     "mean",
     "meta_namespace",
     "pickle_without",
@@ -316,48 +322,39 @@ def capabilities(xp: ModuleType) -> dict[str, int]:
     return out
 
 
-# Helper of ``extract_objects`` and ``repack_objects``
-_repacking_objects: ContextVar[Iterator[object]] = ContextVar("_repacking_objects")
-
-
-def _expand() -> object:  # numpydoc ignore=RT01
-    """
-    Helper of ``extract_objects`` and ``repack_objects``.
-
-    Inverse of the reducer function.
-
-    Notes
-    -----
-    This function must be global in order to be picklable.
-    """
-    try:
-        return next(_repacking_objects.get())
-    except StopIteration:
-        msg = "Not enough objects to repack"
-        raise ValueError(msg)
+_BASIC_TYPES = frozenset((
+    NoneType, bool, int, float, complex, str, bytes, bytearray,
+    list, tuple, dict, set, frozenset, range, slice,
+))  # fmt: skip
 
 
-def pickle_without(obj: object, *classes: type[T]) -> tuple[bytes, list[T]]:
+def pickle_without(
+    obj: object, cls: type[T] | tuple[type[T], ...] = ()
+) -> tuple[bytes, tuple[T, ...], tuple[object, ...]]:
     """
-    Variant of ``pickle.dumps`` that extracts inner objects.
+    Variant of ``pickle.dumps`` that always succeeds and extracts inner objects.
 
     Conceptually, this is similar to passing the ``buffer_callback`` argument to
-    ``pickle.dumps``, but instead of extracting buffers it extracts entire objects.
+    ``pickle.dumps``, but instead of extracting buffers it extracts entire objects,
+    which are either not serializable with ``pickle`` (e.g. local classes or functions)
+    or instances of an explicit list of classes.
 
     Parameters
     ----------
     obj : object
         The object to pickle.
-    *classes : type
-        One or more classes to extract from the object.
+    cls : type | tuple[type, ...], optional
+        One or multiple classes to extract from the object.
         The instances of these classes inside ``obj`` will not be pickled.
 
     Returns
     -------
     bytes
         The pickled object. Must be unpickled with :func:`unpickle_without`.
-    list
-        All instances of ``classes`` found inside ``obj`` (not pickled).
+    tuple
+        All instances of ``cls`` found inside ``obj`` (not pickled).
+    tuple
+        All other objects which failed to pickle.
 
     See Also
     --------
@@ -366,75 +363,221 @@ def pickle_without(obj: object, *classes: type[T]) -> tuple[bytes, list[T]]:
 
     Examples
     --------
+    >>> class NS:
+    ...     def __repr__(self):
+    ...         return "<NS>"
+    ...     def __reduce__(self):
+    ...         assert False
     >>> class A:
     ...     def __repr__(self):
     ...         return "<A>"
-    ...     def __reduce__(self):
-    ...         assert False, "Not serializable"
-    >>> obj = {1: A(), 2: [A(), A()]}  # Any serializable object
-    >>> pik, extracted = pickle_without(obj, A)
-    >>> extracted
-    [<A>, <A>, <A>]
-    >>> unpickle_without(pik, extracted)
-    {1: <A>, 2: [<A>, <A>]}
+    >>> obj = {1: A(), 2: [A(), NS(), A()]}
+    >>> pik, instances, unpickleable = pickle_without(obj, A)
+    >>> instances, unpickleable
+    ([<A>, <A>, <A>], [<NS>])
+    >>> unpickle_without(pik, instances, unpickleable)
+    {1: <A>, 2: [<A>, <NS>, <A>]}
 
     This can be also used to hot-swap inner objects; the only constraint is that
     the number of objects in and out must be the same:
 
     >>> class B:
     ...     def __repr__(self): return "<B>"
-    >>> unpickle_without(pik, [B(), B(), B()])
-    {1: <B>, 2: [<B>, <B>]}
+    >>> unpickle_without(pik, [B(), B(), B()], [NS()])
+    {1: <B>, 2: [<B>, <NS>, <B>]}
     """
-    extracted = []
-
-    def reduce(x: T) -> tuple[Callable[[], object], tuple[()]]:  # numpydoc ignore=GL08
-        extracted.append(x)
-        return _expand, ()
+    instances: list[T] = []
+    unpickleable: list[object] = []
+    seen: dict[int, Literal[0, 1, None]] = {}
+
+    class Pickler(pickle.Pickler):  # numpydoc ignore=GL01,RT01
+        """Override pickle.Pickler.persistent_id.
+
+        TODO consider moving to top-level scope to allow for
+        the full Pickler API to be used.
+        """
+
+        @override
+        def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
+            # Fast exit in case of basic builtin types.
+            # Note that basic collections (tuple, list, etc.) are in this;
+            # persistent_id() will be called again with their contents.
+            if type(obj) in _BASIC_TYPES:  # No subclasses!
+                return None
+
+            id_ = id(obj)
+            try:
+                kind = seen[id_]
+                return None if kind is None else (id_, kind)
+            except KeyError:
+                pass
+
+            if isinstance(obj, cls):
+                instances.append(obj)  # type: ignore[arg-type]
+                seen[id_] = 0
+                return id_, 0
+
+            for func in (
+                # Note: a class that defines __slots__ without defining __getstate__
+                # cannot be pickled with __reduce__(), but can with __reduce_ex__(5)
+                lambda: obj.__reduce_ex__(pickle.HIGHEST_PROTOCOL),
+                lambda: obj.__reduce__(),
+                # Global functions don't have __reduce__, which can be pickled
+                lambda: pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL),
+            ):
+                try:
+                    # a class that defines __slots__ without defining __getstate__
+                    # cannot be pickled with __reduce__(), but can with __reduce_ex__(5)
+                    func()
+                except Exception:  # pylint: disable=broad-exception-caught
+                    pass
+                else:  # Can be pickled
+                    seen[id_] = None
+                    return None
+
+            # Can't be pickled
+            unpickleable.append(obj)
+            seen[id_] = 1
+            return id_, 1
 
     f = io.BytesIO()
-    p = pickle.Pickler(f)
-
-    # Override the reducer for the given classes and all their
-    # subclasses (recursively).
-    p.dispatch_table = copyreg.dispatch_table.copy()
-    subclasses = list(classes)
-    while subclasses:
-        cls = subclasses.pop()
-        p.dispatch_table[cls] = reduce
-        subclasses.extend(cls.__subclasses__())
-
+    p = Pickler(f, protocol=pickle.HIGHEST_PROTOCOL)
     p.dump(obj)
+    return f.getvalue(), tuple(instances), tuple(unpickleable)
 
-    return f.getvalue(), extracted
 
-
-def unpickle_without(pik: bytes, objects: Iterable[object], /) -> Any:  # type: ignore[explicit-any]
+def unpickle_without(  # type: ignore[explicit-any]
+    pik: bytes,
+    instances: Iterable[object],
+    unpickleable: Iterable[object],
+    /,
+) -> Any:
     """
     Variant of ``pickle.loads``, reverse of ``pickle_without``.
 
     Parameters
     ----------
     pik : bytes
         The pickled object generated by ``pickle_without``.
-    objects : Iterable
-        The objects to be reinserted into the unpickled object.
-        Must be the at least the same number of elements as the ones extracted by
-        ``pickle_without``, but does not need to be the same objects or even the
-        same types of objects. Excess objects, if any, won't be inserted.
+    instances : Iterable[object]
+        Instances of the class or classes explicitly passed to ``pickle_without``,
+        to be reinserted into the unpickled object.
+    unpickleable : Iterable[object]
+        The objects that failed to pickle, as returned by ``pickle_without``.
 
     Returns
     -------
     object
-        The unpickled object, with the objects in ``objects`` inserted back into it.
+        The unpickled object.
 
     See Also
     --------
     pickle_without : Serializing function.
     pickle.loads : Standard unpickle function.
+
+    Notes
+    -----
+    The second and third parameter of this function must yield at least the same number
+    of elements as the ones returned by ``pickle_without``, but do not need to be the
+    same objects, or even the same types of objects. Excess objects, if any, will be
+    quietly ignored.
+    """
+    iters = iter(instances), iter(unpickleable)
+    seen: dict[tuple[int, int], object] = {}
+
+    class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL01,RT01
+        """
+        Override pickle.Pickler.persistent_load.
+
+        TODO consider moving to top-level scope to allow for
+        the full Unpickler API to be used.
+        """
+
+        @override
+        def persistent_load(self, pid: tuple[int, int]) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # numpydoc ignore=GL08
+            try:
+                return seen[pid]
+            except KeyError:
+                pass
+
+            _, kind = pid
+            try:
+                obj = next(iters[kind])
+            except StopIteration as e:
+                msg = "Not enough objects to unpickle"
+                raise ValueError(msg) from e
+
+            seen[pid] = obj
+            return obj
+
+    f = io.BytesIO(pik)
+    return Unpickler(f).load()
+
+
+def jax_autojit(
+    func: Callable[P, T],
+) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01,SS03
+    """
+    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``, static arguments must be either hashable or serializable with
+      ``pickle``.
+    - 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, then rebuild the arguments when entering `func`, swapping the JAX
+      concrete arrays with tracer objects.
+    - 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.
     """
-    tok = _repacking_objects.set(iter(objects))
-    try:
-        return pickle.loads(pik)
-    finally:
-        _repacking_objects.reset(tok)
+    import jax
+
+    # {
+    #   jit_cache_key(args_pik, args_arrays, args_unpickleable):
+    #   (res_pik, res_unpickleable)
+    # }
+    static_return_values: dict[Hashable, tuple[bytes, tuple[object, ...]]] = {}
+
+    def jit_cache_key(  # type: ignore[no-any-unimported]  # numpydoc ignore=GL08
+        args_pik: bytes,
+        args_arrays: tuple[jax.Array, ...],  # pyright: ignore[reportUnknownParameterType]
+        args_unpickleable: tuple[Hashable, ...],
+    ) -> Hashable:
+        return (
+            args_pik,
+            tuple((arr.shape, arr.dtype) for arr in args_arrays),  # pyright: ignore[reportUnknownArgumentType]
+            args_unpickleable,
+        )
+
+    def inner(  # type: ignore[no-any-unimported]  # pyright: ignore[reportUnknownParameterType]
+        args_pik: bytes,
+        args_arrays: tuple[jax.Array, ...],  # pyright: ignore[reportUnknownParameterType]
+        args_unpickleable: tuple[Hashable, ...],
+    ) -> tuple[jax.Array, ...]:  # numpydoc ignore=GL08
+        args, kwargs = unpickle_without(args_pik, args_arrays, args_unpickleable)  # pyright: ignore[reportUnknownArgumentType]
+        res = func(*args, **kwargs)  # pyright: ignore[reportCallIssue]
+        res_pik, res_arrays, res_unpickleable = pickle_without(res, jax.Array)  # pyright: ignore[reportUnknownArgumentType]
+        key = jit_cache_key(args_pik, args_arrays, args_unpickleable)
+        val = res_pik, res_unpickleable
+        prev = static_return_values.setdefault(key, val)
+        assert prev == val, "cache key collision"
+        return res_arrays
+
+    jitted = jax.jit(inner, static_argnums=(0, 2))
+
+    @wraps(func)
+    def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
+        args_pik, args_arrays, args_unpickleable = pickle_without(
+            (args, kwargs),
+            jax.Array,  # pyright: ignore[reportUnknownArgumentType]
+        )
+        res_arrays = jitted(args_pik, args_arrays, args_unpickleable)
+        key = jit_cache_key(args_pik, args_arrays, args_unpickleable)
+        res_pik, res_unpickleable = static_return_values[key]
+        return unpickle_without(res_pik, res_arrays, res_unpickleable)  # pyright: ignore[reportUnknownArgumentType]
+
+    return outer