light version

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+import copyreg
+import io
 import math
-from collections.abc import Generator, Iterable
+import pickle
+from collections.abc import Callable, Generator, Iterable, Iterator
+from contextvars import ContextVar
 from types import ModuleType
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 from . import _compat
 from ._compat import (
@@ -22,6 +26,8 @@
     # TODO import from typing (requires Python >=3.13)
     from typing_extensions import TypeIs
 
+T = TypeVar("T")
+
 
 __all__ = [
     "asarrays",
@@ -31,6 +37,8 @@
     "is_python_scalar",
     "mean",
     "meta_namespace",
+    "pickle_without",
+    "unpickle_without",
 ]
 
 
@@ -306,3 +314,116 @@ def capabilities(xp: ModuleType) -> dict[str, int]:
         out["boolean indexing"] = True
         out["data-dependent shapes"] = True
     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.
+    """
+    return next(_repacking_objects.get())
+
+
+def pickle_without(obj: object, *classes: type[T]) -> tuple[bytes, list[T]]:
+    """
+    Variant of ``pickle.dumps`` that extracts inner objects.
+
+    Conceptually, this is similar as passing the ``buffer_callback`` argument to
+    ``pickle.dumps``, but instead of extracting buffers it extracts entire objects.
+
+    Parameters
+    ----------
+    obj : object
+        The object to pickle.
+    *classes : type
+        One or more 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).
+
+    See Also
+    --------
+    pickle.dumps : Standard pickle function.
+    unpickle_without : Reverse function.
+
+    Examples
+    --------
+    >>> 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>]}
+
+    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>]}
+    """
+    extracted = []
+
+    def reduce(x: T) -> tuple[Callable[[], object], tuple[()]]:  # numpydoc ignore=GL08
+        extracted.append(x)
+        return _expand, ()
+
+    f = io.BytesIO()
+    p = pickle.Pickler(f)
+    p.dispatch_table = copyreg.dispatch_table.copy()
+    for cls in classes:
+        p.dispatch_table[cls] = reduce
+    p.dump(obj)
+
+    return f.getvalue(), extracted
+
+
+def unpickle_without(pik: bytes, objects: Iterable[object], /) -> Any:  # type: ignore[explicit-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.
+
+    Returns
+    -------
+    object
+        The unpickled object, with the objects in ``objects`` inserted back into it.
+
+    See Also
+    --------
+    pickle_without : Serializing function.
+    pickle.loads : Standard unpickle function.
+    """
+    tok = _repacking_objects.set(iter(objects))
+    try:
+        return pickle.loads(pik)
+    finally:
+        _repacking_objects.reset(tok)

src/array_api_extra/testing.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
 
 from ._lib._utils._compat import is_dask_namespace, is_jax_namespace
+from ._lib._utils._helpers import pickle_without, unpickle_without
 
 __all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
 
@@ -243,14 +244,10 @@ def iter_tagged() -> (  # type: ignore[explicit-any]
 
         for mod, name, func, tags in iter_tagged():
             if tags["jax_jit"]:
-                # suppress unused-ignore to run mypy in -e lint as well as -e dev
-                wrapped = cast(  # type: ignore[explicit-any]
-                    Callable[..., Any],
-                    jax.jit(
-                        func,
-                        static_argnums=tags["static_argnums"],
-                        static_argnames=tags["static_argnames"],
-                    ),
+                wrapped = _jax_wrap(
+                    func,
+                    static_argnums=tags["static_argnums"],
+                    static_argnames=tags["static_argnames"],
                 )
                 monkeypatch.setattr(mod, name, wrapped)
 
@@ -300,6 +297,7 @@ def _dask_wrap(
     After the function returns, materialize the graph in order to re-raise exceptions.
     """
     import dask
+    import dask.array as da
 
     func_name = getattr(func, "__name__", str(func))
     n_str = f"only up to {n}" if n else "no"
@@ -319,6 +317,40 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         # Block until the graph materializes and reraise exceptions. This allows
         # `pytest.raises` and `pytest.warns` to work as expected. Note that this would
         # not work on scheduler='distributed', as it would not block.
-        return dask.persist(out, scheduler="threads")[0]  # type: ignore[attr-defined,no-untyped-call,func-returns-value,index]  # pyright: ignore[reportPrivateImportUsage]
+        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 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
+    """
+    Wrap `func` inside ``jax.jit``.
+
+    Accepts non-array return values.
+    """
+    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
+
+    jitted = jax.jit(
+        inner, static_argnums=static_argnums, static_argnames=static_argnames
+    )
+    cpu = jax.devices("cpu")[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:])
+
+    return outer