refactor

Author: crusaderky

src/array_api_extra/_lib/_utils/_helpers.py

@@ -5,9 +5,10 @@
 import io
 import math
 import pickle
-from collections.abc import Generator, Iterable
+from collections.abc import Callable, Generator, Hashable, Iterable
+from functools import wraps
 from types import ModuleType, NoneType
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Literal, ParamSpec, TypeVar, cast
 
 from . import _compat
 from ._compat import (
@@ -29,6 +30,7 @@ def override(func):
         return func
 
 
+P = ParamSpec("P")
 T = TypeVar("T")
 
 
@@ -38,6 +40,7 @@ def override(func):
     "eager_shape",
     "in1d",
     "is_python_scalar",
+    "jax_autojit",
     "mean",
     "meta_namespace",
     "pickle_without",
@@ -368,7 +371,7 @@ def pickle_without(
     >>> class A:
     ...     def __repr__(self):
     ...         return "<A>"
-    >>> obj = {1: A(), 2: [A(), NS(), A()]}  # Any serializable object
+    >>> obj = {1: A(), 2: [A(), NS(), A()]}
     >>> pik, instances, unpickleable = pickle_without(obj, A)
     >>> instances, unpickleable
     ([<A>, <A>, <A>], [<NS>])
@@ -396,7 +399,6 @@ class Pickler(pickle.Pickler):  # numpydoc ignore=GL01,RT01
 
         @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.
@@ -416,7 +418,9 @@ def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompa
                 return id_, 0
 
             try:
-                _ = obj.__reduce__()
+                # a class that defines __slots__ without defining __getstate__
+                # cannot be pickled with __reduce__(), but can with __reduce_ex__(5)
+                _ = obj.__reduce_ex__(pickle.HIGHEST_PROTOCOL)
             except Exception:  # pylint: disable=broad-exception-caught
                 pass
             else:  # Can be pickled
@@ -425,7 +429,7 @@ def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompa
 
             # May be a global function, which can be pickled
             try:
-                _ = pickle.dumps(obj)
+                _ = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
             except Exception:  # pylint: disable=broad-exception-caught
                 pass
             else:  # Can be pickled
@@ -438,7 +442,7 @@ def persistent_id(self, obj: object) -> object:  # pyright: ignore[reportIncompa
             return id_, 1
 
     f = io.BytesIO()
-    p = Pickler(f)
+    p = Pickler(f, protocol=pickle.HIGHEST_PROTOCOL)
     p.dump(obj)
     return f.getvalue(), tuple(instances), tuple(unpickleable)
 
@@ -480,7 +484,7 @@ def unpickle_without(  # type: ignore[explicit-any]
     quietly ignored.
     """
     iters = iter(instances), iter(unpickleable)
-    seen: dict[int, object] = {}
+    seen: dict[tuple[int, int], object] = {}
 
     class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL01,RT01
         """
@@ -509,3 +513,72 @@ def persistent_load(self, pid: tuple[int, int]) -> object:  # pyright: ignore[re
 
     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.
+    """
+    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

src/array_api_extra/testing.py

@@ -9,13 +9,13 @@
 import contextlib
 import enum
 import warnings
-from collections.abc import Callable, Hashable, Iterator, Sequence
+from collections.abc import Callable, Iterator, Sequence
 from functools import wraps
 from types import ModuleType
 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
+from ._lib._utils._helpers import jax_autojit, pickle_without, unpickle_without
 
 __all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
 
@@ -29,7 +29,7 @@
     # Sphinx hacks
     SchedulerGetCallable = object
 
-    def override(func: object) -> object:
+    def override(func):
         return func
 
 
@@ -260,7 +260,7 @@ def iter_tagged() -> (  # type: ignore[explicit-any]
     elif is_jax_namespace(xp):
         for mod, name, func, tags in iter_tagged():
             if tags["jax_jit"]:
-                wrapped = _jax_autojit(func)
+                wrapped = jax_autojit(func)
                 monkeypatch.setattr(mod, name, wrapped)
 
 
@@ -334,66 +334,3 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         return unpickle_without(pik, arrays, unpickleable)  # pyright: ignore[reportUnknownArgumentType]
 
     return wrapper
-
-
-def _jax_autojit(func: Callable[P, T]) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01
-    """
-    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.
-    - 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
-
-    # pickled return values of `func`, minus the JAX arrays
-    static_return_values: dict[Hashable, tuple[bytes, tuple[object, ...]]] = {}
-
-    def jit_cache_key(  # type: ignore[no-any-unimported]
-        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