, "baz"]}
+ """
+ instances: list[T] = []
+ rest: list[object] = []
+
+ class Pickler(pickle.Pickler): # numpydoc ignore=GL08
+ """
+ Use the `pickle.Pickler.persistent_id` hook to extract objects.
+ """
+
+ @override
+ def persistent_id(
+ self, obj: object
+ ) -> Literal[0, 1, None]: # numpydoc ignore=GL08
+ if isinstance(obj, cls):
+ instances.append(obj) # type: ignore[arg-type]
+ return 0
+
+ typ_ = type(obj)
+ if typ_ in _BASIC_PICKLED_TYPES: # No subclasses!
+ # If obj is a collection, recursively descend inside it
+ return None
+ if typ_ in _BASIC_REST_TYPES:
+ rest.append(obj)
+ return 1
+
+ try:
+ # Note: 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
+ rest.append(obj)
+ return 1
+
+ # Object can be pickled. Let the Pickler recursively descend inside it.
+ return None
+
+ f = io.BytesIO()
+ p = Pickler(f, protocol=pickle.HIGHEST_PROTOCOL)
+ p.dump(obj)
+ return instances, (f.getvalue(), *rest)
+
+
+def pickle_unflatten(instances: Iterable[object], rest: FlattenRest) -> Any: # type: ignore[explicit-any]
+ """
+ Reverse of ``pickle_flatten``.
+
+ Parameters
+ ----------
+ instances : Iterable
+ Inner objects to be reinserted into the flattened container.
+ rest : FlattenRest
+ Extra bits, as returned by ``pickle_flatten``.
+
+ Returns
+ -------
+ object
+ The outer object originally passed to ``pickle_flatten`` after a
+ pickle->unpickle round-trip.
+
+ See Also
+ --------
+ pickle_flatten : Serializing function.
+ pickle.loads : Standard unpickle function.
+
+ Notes
+ -----
+ The `instances` iterable must yield at least the same number of elements as the ones
+ returned by ``pickle_flatten``, but the elements do not need to be the same objects
+ or even the same types of objects. Excess elements, if any, will be left untouched.
+ """
+ iters = iter(instances), iter(rest)
+ pik = cast(bytes, next(iters[1]))
+
+ class Unpickler(pickle.Unpickler): # numpydoc ignore=GL08
+ """Mirror of the overridden Pickler in pickle_flatten."""
+
+ @override
+ def persistent_load(self, pid: Literal[0, 1]) -> object: # numpydoc ignore=GL08
+ try:
+ return next(iters[pid])
+ except StopIteration as e:
+ msg = "Not enough objects to unpickle"
+ raise ValueError(msg) from e
+
+ f = io.BytesIO(pik)
+ return Unpickler(f).load()
+
+
+class _AutoJITWrapper(Generic[T]): # numpydoc ignore=PR01
+ """
+ Helper of :func:`jax_autojit`.
+
+ Wrap arbitrary inputs and outputs of the jitted function and
+ convert them to/from PyTrees.
+ """
+
+ obj: T
+ _registered: ClassVar[bool] = False
+ __slots__: tuple[str, ...] = ("obj",)
+
+ def __init__(self, obj: T) -> None: # numpydoc ignore=GL08
+ self._register()
+ self.obj = obj
+
+ @classmethod
+ def _register(cls): # numpydoc ignore=SS06
+ """
+ Register upon first use instead of at import time, to avoid
+ globally importing JAX.
+ """
+ if not cls._registered:
+ import jax
+
+ jax.tree_util.register_pytree_node(
+ cls,
+ lambda obj: pickle_flatten(obj, jax.Array), # pyright: ignore[reportUnknownArgumentType]
+ lambda aux_data, children: pickle_unflatten(children, aux_data), # pyright: ignore[reportUnknownArgumentType]
+ )
+ cls._registered = True
+
+
+def jax_autojit(
+ func: Callable[P, T],
+) -> Callable[P, T]: # numpydoc ignore=PR01,RT01,SS03
+ """
+ Wrap `func` with ``jax.jit``, with the following differences:
+
+ - Python scalar 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.
+
+ See Also
+ --------
+ jax.jit : JAX JIT compilation function.
+
+ Notes
+ -----
+ These are useful choices *for testing purposes only*, which is how this function is
+ intended to be used. The output of ``jax.jit`` is a C++ level callable, that
+ directly dispatches to the compiled kernel after the initial call. In comparison,
+ ``jax_autojit`` incurs a much higher dispatch time.
+
+ Additionally, consider::
+
+ def f(x: Array, y: float, plus: bool) -> Array:
+ return x + y if plus else x - y
+
+ j1 = jax.jit(f, static_argnames="plus")
+ j2 = jax_autojit(f)
+
+ In the above example, ``j2`` requires a lot less setup to be tested effectively than
+ ``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
+
+ @jax.jit # type: ignore[misc] # pyright: ignore[reportUntypedFunctionDecorator]
+ def inner( # type: ignore[decorated-any,explicit-any] # numpydoc ignore=GL08
+ wargs: _AutoJITWrapper[Any],
+ ) -> _AutoJITWrapper[T]:
+ args, kwargs = wargs.obj
+ res = func(*args, **kwargs) # pyright: ignore[reportCallIssue]
+ return _AutoJITWrapper(res)
+
+ @wraps(func)
+ def outer(*args: P.args, **kwargs: P.kwargs) -> T: # numpydoc ignore=GL08
+ wargs = _AutoJITWrapper((args, kwargs))
+ return inner(wargs).obj
+
+ return outer
diff --git a/src/array_api_extra/_lib/_utils/_typing.py b/src/array_api_extra/_lib/_utils/_typing.py
index d32a3a07..8204be47 100644
--- a/src/array_api_extra/_lib/_utils/_typing.py
+++ b/src/array_api_extra/_lib/_utils/_typing.py
@@ -1,5 +1,5 @@
# numpydoc ignore=GL08
-# pylint: disable=missing-module-docstring
+# pylint: disable=missing-module-docstring,duplicate-code
Array = object
DType = object
diff --git a/src/array_api_extra/testing.py b/src/array_api_extra/testing.py
index 37e8e69e..3979f9dd 100644
--- a/src/array_api_extra/testing.py
+++ b/src/array_api_extra/testing.py
@@ -7,12 +7,15 @@
from __future__ import annotations
import contextlib
-from collections.abc import Callable, Iterable, Iterator, Sequence
+import enum
+import warnings
+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 jax_autojit, pickle_flatten, pickle_unflatten
__all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
@@ -26,7 +29,7 @@
# Sphinx hacks
SchedulerGetCallable = object
- def override(func: object) -> object:
+ def override(func):
return func
@@ -36,13 +39,22 @@ def override(func: object) -> object:
_ufuncs_tags: dict[object, dict[str, Any]] = {} # type: ignore[explicit-any]
+class Deprecated(enum.Enum):
+ """Unique type for deprecated parameters."""
+
+ DEPRECATED = 1
+
+
+DEPRECATED = Deprecated.DEPRECATED
+
+
def lazy_xp_function( # type: ignore[explicit-any]
func: Callable[..., Any],
*,
allow_dask_compute: bool | int = False,
jax_jit: bool = True,
- static_argnums: int | Sequence[int] | None = None,
- static_argnames: str | Iterable[str] | None = None,
+ static_argnums: Deprecated = DEPRECATED,
+ static_argnames: Deprecated = DEPRECATED,
) -> None: # numpydoc ignore=GL07
"""
Tag a function to be tested on lazy backends.
@@ -82,16 +94,31 @@ def lazy_xp_function( # type: ignore[explicit-any]
Default: False, meaning that `func` must be fully lazy and never materialize the
graph.
jax_jit : bool, optional
- Set to True to replace `func` with ``jax.jit(func)`` after calling the
- :func:`patch_lazy_xp_functions` test helper with ``xp=jax.numpy``. Set to False
- if `func` is only compatible with eager (non-jitted) JAX. Default: True.
- static_argnums : int | Sequence[int], optional
- Passed to jax.jit. Positional arguments to treat as static (compile-time
- constant). Default: infer from `static_argnames` using
- `inspect.signature(func)`.
- static_argnames : str | Iterable[str], optional
- Passed to jax.jit. Named arguments to treat as static (compile-time constant).
- Default: infer from `static_argnums` using `inspect.signature(func)`.
+ Set to True to replace `func` with a smart variant of ``jax.jit(func)`` after
+ calling the :func:`patch_lazy_xp_functions` test helper with ``xp=jax.numpy``.
+ This is the default behaviour.
+ Set to False if `func` is only compatible with eager (non-jitted) JAX.
+
+ Unlike with vanilla ``jax.jit``, all arguments and return types that are not JAX
+ arrays are treated as static; the function can accept and return arbitrary
+ wrappers around JAX arrays. This difference is because, in real life, most users
+ won't wrap the function directly with ``jax.jit`` but rather they will use it
+ within their own code, which is itself then wrapped by ``jax.jit``, and
+ internally consume the function's outputs.
+
+ In other words, the pattern that is being tested is::
+
+ >>> @jax.jit
+ ... def user_func(x):
+ ... y = user_prepares_inputs(x)
+ ... z = func(y, some_static_arg=True)
+ ... return user_consumes(z)
+
+ Default: True.
+ static_argnums :
+ Deprecated; ignored
+ static_argnames :
+ Deprecated; ignored
See Also
--------
@@ -108,7 +135,7 @@ def lazy_xp_function( # type: ignore[explicit-any]
def test_myfunc(xp):
a = xp.asarray([1, 2])
- # When xp=jax.numpy, this is the same as `b = jax.jit(myfunc)(a)`
+ # When xp=jax.numpy, this is similar to `b = jax.jit(myfunc)(a)`
# When xp=dask.array, crash on compute() or persist()
b = myfunc(a)
@@ -168,12 +195,20 @@ def test_myfunc(xp):
b = mymodule.myfunc(a) # This is wrapped when xp=jax.numpy or xp=dask.array
c = naked.myfunc(a) # This is not
"""
+ if static_argnums is not DEPRECATED or static_argnames is not DEPRECATED:
+ warnings.warn(
+ (
+ "The `static_argnums` and `static_argnames` parameters are deprecated "
+ "and ignored. They will be removed in a future version."
+ ),
+ DeprecationWarning,
+ stacklevel=2,
+ )
tags = {
"allow_dask_compute": allow_dask_compute,
"jax_jit": jax_jit,
- "static_argnums": static_argnums,
- "static_argnames": static_argnames,
}
+
try:
func._lazy_xp_function = tags # type: ignore[attr-defined] # pylint: disable=protected-access # pyright: ignore[reportFunctionMemberAccess]
except AttributeError: # @cython.vectorize
@@ -247,19 +282,9 @@ def iter_tagged() -> ( # type: ignore[explicit-any]
monkeypatch.setattr(mod, name, wrapped)
elif is_jax_namespace(xp):
- import jax
-
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_autojit(func)
monkeypatch.setattr(mod, name, wrapped)
@@ -308,6 +333,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"
@@ -327,6 +353,8 @@ 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]
+ arrays, rest = pickle_flatten(out, da.Array)
+ arrays = dask.persist(arrays, scheduler="threads")[0] # type: ignore[attr-defined,no-untyped-call,func-returns-value,index] # pyright: ignore[reportPrivateImportUsage]
+ return pickle_unflatten(arrays, rest) # pyright: ignore[reportUnknownArgumentType]
return wrapper
diff --git a/tests/conftest.py b/tests/conftest.py
index 5676cc0d..82a0acbb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -148,31 +148,45 @@ def xp(
patch_lazy_xp_functions(request, monkeypatch, xp=xp)
if library.like(Backend.JAX):
- import jax
+ _setup_jax(library)
+ elif library.like(Backend.TORCH):
+ _setup_torch(library)
- # suppress unused-ignore to run mypy in -e lint as well as -e dev
- jax.config.update("jax_enable_x64", True) # type: ignore[no-untyped-call,unused-ignore]
+ yield xp
- if library == Backend.JAX_GPU:
- try:
- device = jax.devices("cuda")[0]
- except RuntimeError:
- pytest.skip("no CUDA device available")
- else:
- device = jax.devices("cpu")[0]
- jax.config.update("jax_default_device", device)
-
- elif library == Backend.TORCH_GPU:
- import torch.cuda
- if not torch.cuda.is_available():
+def _setup_jax(library: Backend) -> None:
+ import jax
+
+ # suppress unused-ignore to run mypy in -e lint as well as -e dev
+ jax.config.update("jax_enable_x64", True) # type: ignore[no-untyped-call,unused-ignore]
+
+ if library == Backend.JAX_GPU:
+ try:
+ device = jax.devices("cuda")[0]
+ except RuntimeError:
pytest.skip("no CUDA device available")
- xp.set_default_device("cuda")
+ else:
+ device = jax.devices("cpu")[0]
+ jax.config.update("jax_default_device", device)
- elif library == Backend.TORCH: # CPU
- xp.set_default_device("cpu")
- yield xp
+def _setup_torch(library: Backend) -> None:
+ import torch
+
+ # This is already the default, but some tests or env variables may change it.
+ # TODO test both float32 and float64, like in scipy.
+ torch.set_default_dtype(torch.float32)
+
+ if library == Backend.TORCH_GPU:
+ import torch.cuda
+
+ if not torch.cuda.is_available():
+ pytest.skip("no CUDA device available")
+ torch.set_default_device("cuda")
+ else:
+ assert library == Backend.TORCH
+ torch.set_default_device("cpu")
@pytest.fixture(params=[Backend.DASK]) # Can select the test with `pytest -k dask`
@@ -186,6 +200,26 @@ def da(
return xp
+@pytest.fixture(params=[Backend.JAX, Backend.JAX_GPU])
+def jnp(
+ request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> ModuleType: # numpydoc ignore=PR01,RT01
+ """Variant of the `xp` fixture that only yields jax.numpy."""
+ xp = pytest.importorskip("jax.numpy")
+ _setup_jax(request.param)
+ patch_lazy_xp_functions(request, monkeypatch, xp=xp)
+ return xp
+
+
+@pytest.fixture(params=[Backend.TORCH, Backend.TORCH_GPU])
+def torch(request: pytest.FixtureRequest) -> ModuleType: # numpydoc ignore=PR01,RT01
+ """Variant of the `xp` fixture that only yields torch."""
+ xp = pytest.importorskip("torch")
+ xp = array_namespace(xp.empty(0))
+ _setup_torch(request.param)
+ return xp
+
+
@pytest.fixture
def device(
library: Backend, xp: ModuleType
@@ -196,7 +230,9 @@ def device(
Where possible, return a device that is not the default one.
"""
if library == Backend.ARRAY_API_STRICT:
- d = xp.Device("device1")
- assert get_device(xp.empty(0)) != d
- return d
+ return xp.Device("device1")
+ if library == Backend.TORCH:
+ return xp.device("meta")
+ if library == Backend.TORCH_GPU:
+ return xp.device("cpu")
return get_device(xp.empty(0))
diff --git a/tests/test_at.py b/tests/test_at.py
index fa9bcdc8..7294a7c4 100644
--- a/tests/test_at.py
+++ b/tests/test_at.py
@@ -1,5 +1,4 @@
import math
-import pickle
from collections.abc import Callable, Generator
from contextlib import contextmanager
from types import ModuleType
@@ -41,28 +40,11 @@ def at_op(
just a workaround for when one wants to apply jax.jit to `at()` directly,
which is not a common use case.
"""
- if isinstance(idx, (slice | tuple)):
- return _at_op(x, None, pickle.dumps(idx), op, y, copy=copy, xp=xp)
- return _at_op(x, idx, None, op, y, copy=copy, xp=xp)
-
-
-def _at_op(
- x: Array,
- idx: SetIndex | None,
- idx_pickle: bytes | None,
- op: _AtOp,
- y: Array | object,
- copy: bool | None,
- xp: ModuleType | None = None,
-) -> Array:
- """jitted helper of at_op"""
- if idx_pickle:
- idx = pickle.loads(idx_pickle)
- meth = cast(Callable[..., Array], getattr(at(x, cast(SetIndex, idx)), op.value)) # type: ignore[explicit-any]
+ meth = cast(Callable[..., Array], getattr(at(x, idx), op.value)) # type: ignore[explicit-any]
return meth(y, copy=copy, xp=xp)
-lazy_xp_function(_at_op, static_argnames=("op", "idx_pickle", "copy", "xp"))
+lazy_xp_function(at_op)
@contextmanager
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index 0cee0b4d..99da4a2b 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -17,10 +17,12 @@
broadcast_shapes,
cov,
create_diagonal,
+ default_dtype,
expand_dims,
isclose,
kron,
nunique,
+ one_hot,
pad,
setdiff1d,
sinc,
@@ -37,17 +39,18 @@
# some xp backends are untyped
# mypy: disable-error-code=no-untyped-def
-lazy_xp_function(apply_where, static_argnums=(2, 3), static_argnames="xp")
-lazy_xp_function(atleast_nd, static_argnames=("ndim", "xp"))
-lazy_xp_function(cov, static_argnames="xp")
-lazy_xp_function(create_diagonal, static_argnames=("offset", "xp"))
-lazy_xp_function(expand_dims, static_argnames=("axis", "xp"))
-lazy_xp_function(kron, static_argnames="xp")
-lazy_xp_function(nunique, static_argnames="xp")
-lazy_xp_function(pad, static_argnames=("pad_width", "mode", "constant_values", "xp"))
+lazy_xp_function(apply_where)
+lazy_xp_function(atleast_nd)
+lazy_xp_function(cov)
+lazy_xp_function(create_diagonal)
+lazy_xp_function(expand_dims)
+lazy_xp_function(kron)
+lazy_xp_function(nunique)
+lazy_xp_function(one_hot)
+lazy_xp_function(pad)
# FIXME calls in1d which calls xp.unique_values without size
-lazy_xp_function(setdiff1d, jax_jit=False, static_argnames=("assume_unique", "xp"))
-lazy_xp_function(sinc, static_argnames="xp")
+lazy_xp_function(setdiff1d, jax_jit=False)
+lazy_xp_function(sinc)
class TestApplyWhere:
@@ -196,7 +199,6 @@ def test_device(self, xp: ModuleType, device: Device):
y = apply_where(x % 2 == 0, x, self.f1, fill_value=x)
assert get_device(y) == device
- @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
@pytest.mark.filterwarnings("ignore::RuntimeWarning") # overflows, etc.
@hypothesis.settings(
# The xp and library fixtures are not regenerated between hypothesis iterations
@@ -402,7 +404,6 @@ def test_none(self, args: tuple[tuple[float | None, ...], ...]):
assert actual == expect
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
class TestCov:
def test_basic(self, xp: ModuleType):
xp_assert_close(
@@ -415,6 +416,7 @@ def test_complex(self, xp: ModuleType):
expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128)
xp_assert_close(actual, expect)
+ @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="matmul with nan fillvalue")
def test_empty(self, xp: ModuleType):
with warnings.catch_warnings(record=True):
warnings.simplefilter("always", RuntimeWarning)
@@ -449,6 +451,98 @@ def test_xp(self, xp: ModuleType):
)
+@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="backend doesn't have arange")
+class TestOneHot:
+ @pytest.mark.parametrize("n_dim", range(4))
+ @pytest.mark.parametrize("num_classes", [1, 3, 10])
+ def test_dims_and_classes(self, xp: ModuleType, n_dim: int, num_classes: int):
+ shape = tuple(range(2, 2 + n_dim))
+ rng = np.random.default_rng(2347823)
+ np_x = rng.integers(num_classes, size=shape)
+ x = xp.asarray(np_x)
+ y = one_hot(x, num_classes)
+ assert y.shape == (*x.shape, num_classes)
+ for *i_list, j in ndindex(*shape, num_classes):
+ i = tuple(i_list)
+ assert float(y[(*i, j)]) == (int(x[i]) == j)
+
+ def test_basic(self, xp: ModuleType):
+ actual = one_hot(xp.asarray([0, 1, 2]), 3)
+ expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+ xp_assert_equal(actual, expected)
+
+ actual = one_hot(xp.asarray([1, 2, 0]), 3)
+ expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
+ xp_assert_equal(actual, expected)
+
+ def test_2d(self, xp: ModuleType):
+ actual = one_hot(xp.asarray([[2, 1, 0], [1, 0, 2]]), 3, axis=1)
+ expected = xp.asarray(
+ [
+ [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]],
+ [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
+ ]
+ )
+ xp_assert_equal(actual, expected)
+
+ @pytest.mark.skip_xp_backend(
+ Backend.ARRAY_API_STRICTEST, reason="backend doesn't support Boolean indexing"
+ )
+ def test_abstract_size(self, xp: ModuleType):
+ x = xp.arange(5)
+ x = x[x > 2]
+ actual = one_hot(x, 5)
+ expected = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0]])
+ xp_assert_equal(actual, expected)
+
+ @pytest.mark.skip_xp_backend(
+ Backend.TORCH_GPU, reason="Puts Pytorch into a bad state."
+ )
+ def test_out_of_bound(self, xp: ModuleType):
+ # Undefined behavior. Either return zero, or raise.
+ try:
+ actual = one_hot(xp.asarray([-1, 3]), 3)
+ except IndexError:
+ return
+ expected = xp.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+ xp_assert_equal(actual, expected)
+
+ @pytest.mark.parametrize(
+ "int_dtype",
+ ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"],
+ )
+ def test_int_types(self, xp: ModuleType, int_dtype: str):
+ dtype = getattr(xp, int_dtype)
+ x = xp.asarray([0, 1, 2], dtype=dtype)
+ actual = one_hot(x, 3)
+ expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+ xp_assert_equal(actual, expected)
+
+ def test_custom_dtype(self, xp: ModuleType):
+ actual = one_hot(xp.asarray([0, 1, 2], dtype=xp.int32), 3, dtype=xp.bool)
+ expected = xp.asarray(
+ [[True, False, False], [False, True, False], [False, False, True]]
+ )
+ xp_assert_equal(actual, expected)
+
+ def test_axis(self, xp: ModuleType):
+ expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]).T
+ actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=0)
+ xp_assert_equal(actual, expected)
+
+ actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=-2)
+ xp_assert_equal(actual, expected)
+
+ def test_non_integer(self, xp: ModuleType):
+ with pytest.raises(TypeError):
+ _ = one_hot(xp.asarray([1.0]), 3)
+
+ def test_device(self, xp: ModuleType, device: Device):
+ x = xp.asarray([0, 1, 2], device=device)
+ y = one_hot(x, 3)
+ assert get_device(y) == device
+
+
@pytest.mark.skip_xp_backend(
Backend.SPARSE, reason="read-only backend without .at support"
)
@@ -518,6 +612,38 @@ def test_xp(self, xp: ModuleType):
xp_assert_equal(y, xp.asarray([[1, 0], [0, 2]]))
+class TestDefaultDType:
+ def test_basic(self, xp: ModuleType):
+ assert default_dtype(xp) == xp.empty(0).dtype
+
+ def test_kind(self, xp: ModuleType):
+ assert default_dtype(xp, "real floating") == xp.empty(0).dtype
+ assert default_dtype(xp, "complex floating") == (xp.empty(0) * 1j).dtype
+ assert default_dtype(xp, "integral") == xp.int64
+ assert default_dtype(xp, "indexing") == xp.int64
+
+ with pytest.raises(ValueError, match="Unknown kind"):
+ _ = default_dtype(xp, "foo") # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+
+ def test_device(self, xp: ModuleType, device: Device):
+ # Note: at the moment there are no known namespaces with
+ # device-specific default dtypes.
+ assert default_dtype(xp, device=None) == xp.empty(0).dtype
+ assert default_dtype(xp, device=device) == xp.empty(0).dtype
+
+ def test_torch(self, torch: ModuleType):
+ xp = torch
+ xp.set_default_dtype(xp.float64)
+ assert default_dtype(xp) == xp.float64
+ assert default_dtype(xp, "real floating") == xp.float64
+ assert default_dtype(xp, "complex floating") == xp.complex128
+
+ xp.set_default_dtype(xp.float32)
+ assert default_dtype(xp) == xp.float32
+ assert default_dtype(xp, "real floating") == xp.float32
+ assert default_dtype(xp, "complex floating") == xp.complex64
+
+
class TestExpandDims:
def test_single_axis(self, xp: ModuleType):
"""Trivial case where xpx.expand_dims doesn't add anything to xp.expand_dims"""
@@ -570,7 +696,9 @@ def test_xp(self, xp: ModuleType):
@pytest.mark.filterwarnings( # array_api_strictest
"ignore:invalid value encountered:RuntimeWarning:array_api_strict"
)
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
+@pytest.mark.filterwarnings( # sparse
+ "ignore:invalid value encountered:RuntimeWarning:sparse"
+)
class TestIsClose:
@pytest.mark.parametrize("swap", [False, True])
@pytest.mark.parametrize(
@@ -688,6 +816,7 @@ def test_bool_dtype(self, xp: ModuleType):
isclose(xp.asarray(True), b, atol=1), xp.asarray([True, True, True])
)
+ @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="index by sparse array")
@pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape")
def test_none_shape(self, xp: ModuleType):
a = xp.asarray([1, 5, 0])
@@ -696,6 +825,7 @@ def test_none_shape(self, xp: ModuleType):
a = a[a < 5]
xp_assert_equal(isclose(a, b), xp.asarray([True, False]))
+ @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="index by sparse array")
@pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape")
def test_none_shape_bool(self, xp: ModuleType):
a = xp.asarray([True, True, False])
@@ -732,9 +862,6 @@ def test_device(self, xp: ModuleType, device: Device, equal_nan: bool):
b = xp.asarray([1e-9, 1e-4, xp.nan], device=device)
res = isclose(a, b, equal_nan=equal_nan)
assert get_device(res) == device
- xp_assert_equal(
- isclose(a, b, equal_nan=equal_nan), xp.asarray([True, False, equal_nan])
- )
class TestKron:
@@ -795,7 +922,6 @@ def test_kron_shape(
k = kron(a, b)
assert k.shape == expected_shape
- @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
def test_python_scalar(self, xp: ModuleType):
a = 1
# Test no dtype promotion to xp.asarray(a); use b.dtype
@@ -997,6 +1123,9 @@ def test_all_python_scalars(self, assume_unique: bool):
_ = setdiff1d(0, 0, assume_unique=assume_unique)
@assume_unique
+ @pytest.mark.skip_xp_backend(
+ Backend.TORCH, reason="device='meta' does not support unknown shapes"
+ )
def test_device(self, xp: ModuleType, device: Device, assume_unique: bool):
x1 = xp.asarray([3, 8, 20], device=device)
x2 = xp.asarray([2, 3, 4], device=device)
@@ -1011,8 +1140,8 @@ def test_xp(self, xp: ModuleType):
xp_assert_equal(actual, expected)
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
class TestSinc:
+ @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no linspace")
def test_simple(self, xp: ModuleType):
xp_assert_equal(sinc(xp.asarray(0.0)), xp.asarray(1.0))
w = sinc(xp.linspace(-1, 1, 100))
@@ -1024,6 +1153,7 @@ def test_dtype(self, xp: ModuleType, x: int | complex):
with pytest.raises(ValueError, match="real floating data type"):
_ = sinc(xp.asarray(x))
+ @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange")
def test_3d(self, xp: ModuleType):
x = xp.reshape(xp.arange(18, dtype=xp.float64), (3, 3, 2))
expected = xp.zeros((3, 3, 2), dtype=xp.float64)
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
index a104e93c..74a0ec93 100644
--- a/tests/test_helpers.py
+++ b/tests/test_helpers.py
@@ -1,5 +1,5 @@
from types import ModuleType
-from typing import cast
+from typing import TYPE_CHECKING, Generic, TypeVar, cast
import numpy as np
import pytest
@@ -13,18 +13,31 @@
capabilities,
eager_shape,
in1d,
+ jax_autojit,
meta_namespace,
ndindex,
+ pickle_flatten,
+ pickle_unflatten,
)
from array_api_extra._lib._utils._typing import Array, Device, DType
from array_api_extra.testing import lazy_xp_function
from .conftest import np_compat
+if TYPE_CHECKING:
+ # TODO import from typing (requires Python >=3.12)
+ from typing_extensions import override
+else:
+
+ def override(func):
+ return func
+
# mypy: disable-error-code=no-untyped-usage
+T = TypeVar("T")
+
# FIXME calls xp.unique_values without size
-lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))
+lazy_xp_function(in1d, jax_jit=False)
@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no unique_inverse")
@@ -66,7 +79,6 @@ def test_xp(self, xp: ModuleType):
class TestAsArrays:
- @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
@pytest.mark.parametrize(
("dtype", "b", "defined"),
[
@@ -199,8 +211,210 @@ def test_xp(self, xp: ModuleType):
assert meta_namespace(*args, xp=xp) in (xp, np_compat)
-def test_capabilities(xp: ModuleType):
- expect = {"boolean indexing", "data-dependent shapes"}
- if xp.__array_api_version__ >= "2024.12":
- expect.add("max dimensions")
- assert capabilities(xp).keys() == expect
+class TestCapabilities:
+ def test_basic(self, xp: ModuleType):
+ expect = {"boolean indexing", "data-dependent shapes"}
+ if xp.__array_api_version__ >= "2024.12":
+ expect.add("max dimensions")
+ assert capabilities(xp).keys() == expect
+
+ def test_device(self, xp: ModuleType, library: Backend, device: Device):
+ expect_keys = {"boolean indexing", "data-dependent shapes"}
+ if xp.__array_api_version__ >= "2024.12":
+ expect_keys.add("max dimensions")
+ assert capabilities(xp, device=device).keys() == expect_keys
+
+ if library.like(Backend.TORCH):
+ # The output of capabilities is device-specific.
+
+ # Test that device=None gets the current default device.
+ expect = capabilities(xp, device=device)
+ with xp.device(device):
+ actual = capabilities(xp)
+ assert actual == expect
+
+ # Test that we're accepting anything that is accepted by the
+ # device= parameter in other functions
+ actual = capabilities(xp, device=device.type) # type: ignore[attr-defined] # pyright: ignore[reportUnknownArgumentType,reportAttributeAccessIssue]
+
+
+class Wrapper(Generic[T]):
+ """Trivial opaque wrapper. Must be pickleable."""
+
+ x: T
+ # __slots__ make this object serializable with __reduce_ex__(5),
+ # but not with __reduce__
+ __slots__: tuple[str, ...] = ("x",)
+
+ def __init__(self, x: T):
+ self.x = x
+
+ # Note: this makes the object not hashable
+ @override
+ def __eq__(self, other: object) -> bool:
+ return isinstance(other, Wrapper) and self.x == other.x
+
+
+class TestPickleFlatten:
+ def test_roundtrip(self):
+ class NotSerializable:
+ @override
+ def __reduce__(self) -> tuple[object, ...]:
+ raise NotImplementedError()
+
+ # Note: NotHashable() instances can be reduced to an
+ # unserializable local class
+ class NotHashable:
+ @override
+ def __eq__(self, other: object) -> bool:
+ return isinstance(other, type(self)) and other.__dict__ == self.__dict__
+
+ with pytest.raises(TypeError):
+ _ = hash(NotHashable())
+
+ # Extracted objects need be neither pickleable nor serializable
+ class C(NotSerializable, NotHashable):
+ x: int
+
+ def __init__(self, x: int):
+ self.x = x
+
+ class D(C):
+ pass
+
+ c1 = C(1)
+ c2 = C(2)
+ d3 = D(3)
+
+ # An assorted bunch of opaque containers, standard containers,
+ # non-serializable objects, and non-hashable objects (but not at the same time)
+ obj = Wrapper([1, c1, {2: (c2, {NotSerializable()})}, NotHashable(), d3])
+ instances, rest = pickle_flatten(obj, C)
+
+ assert instances == [c1, c2, d3]
+ obj2 = pickle_unflatten(instances, rest)
+ assert obj2 == obj
+
+ def test_swap_objects(self):
+ class C:
+ pass
+
+ obj = [1, C(), {2: (C(), {C()})}]
+ _, rest = pickle_flatten(obj, C)
+ obj2 = pickle_unflatten(["foo", "bar", "baz"], rest)
+ assert obj2 == [1, "foo", {2: ("bar", {"baz"})}]
+
+ def test_multi_class(self):
+ class C:
+ pass
+
+ class D:
+ pass
+
+ c, d = C(), D()
+ instances, _ = pickle_flatten([c, d], (C, D))
+ assert len(instances) == 2
+ assert instances[0] is c
+ assert instances[1] is d
+
+ def test_no_class(self):
+ obj = {1: "foo", 2: (3, 4)}
+ instances, rest = pickle_flatten(obj, ()) # type: ignore[var-annotated]
+ assert instances == []
+ obj2 = pickle_unflatten([], rest)
+ assert obj2 == obj
+
+ def test_flattened_stream(self):
+ """
+ Test that multiple calls to flatten() can feed into the same stream of instances
+ """
+ obj1 = Wrapper(1)
+ obj2 = [Wrapper(2), Wrapper(3)]
+ instances1, rest1 = pickle_flatten(obj1, Wrapper)
+ instances2, rest2 = pickle_flatten(obj2, Wrapper)
+ it = iter(instances1 + instances2 + [Wrapper(4)]) # pyright: ignore[reportUnknownArgumentType]
+ assert pickle_unflatten(it, rest1) == obj1 # pyright: ignore[reportUnknownArgumentType]
+ assert pickle_unflatten(it, rest2) == obj2 # pyright: ignore[reportUnknownArgumentType]
+ assert list(it) == [Wrapper(4)] # pyright: ignore[reportUnknownArgumentType]
+
+ def test_too_short(self):
+ obj = [Wrapper(1), Wrapper(2)]
+ instances, rest = pickle_flatten(obj, Wrapper)
+ with pytest.raises(ValueError, match="Not enough"):
+ pickle_unflatten(instances[:1], rest) # pyright: ignore[reportUnknownArgumentType]
+
+ def test_recursion(self):
+ obj: list[object] = [Wrapper(1)]
+ obj.append(obj)
+
+ instances, rest = pickle_flatten(obj, Wrapper)
+ assert instances == [Wrapper(1)]
+
+ obj2 = pickle_unflatten(instances, rest) # pyright: ignore[reportUnknownArgumentType]
+ assert len(obj2) == 2
+ assert obj2[0] is obj[0]
+ assert obj2[1] is obj2
+
+
+class TestJAXAutoJIT:
+ def test_basic(self, jnp: ModuleType):
+ @jax_autojit
+ def f(x: Array, k: object = False) -> Array:
+ return x + 1 if k else x - 1
+
+ # Basic recognition of static_argnames
+ xp_assert_equal(f(jnp.asarray([1, 2])), jnp.asarray([0, 1]))
+ xp_assert_equal(f(jnp.asarray([1, 2]), False), jnp.asarray([0, 1]))
+ xp_assert_equal(f(jnp.asarray([1, 2]), True), jnp.asarray([2, 3]))
+ xp_assert_equal(f(jnp.asarray([1, 2]), 1), jnp.asarray([2, 3]))
+
+ # static argument is not an ArrayLike
+ xp_assert_equal(f(jnp.asarray([1, 2]), "foo"), jnp.asarray([2, 3]))
+
+ # static argument is not hashable, but serializable
+ xp_assert_equal(f(jnp.asarray([1, 2]), ["foo"]), jnp.asarray([2, 3]))
+
+ def test_wrapper(self, jnp: ModuleType):
+ @jax_autojit
+ def f(w: Wrapper[Array]) -> Wrapper[Array]:
+ return Wrapper(w.x + 1)
+
+ inp = Wrapper(jnp.asarray([1, 2]))
+ out = f(inp).x
+ xp_assert_equal(out, jnp.asarray([2, 3]))
+
+ def test_static_hashable(self, jnp: ModuleType):
+ """Static argument/return value is hashable, but not serializable"""
+
+ class C:
+ def __reduce__(self) -> object: # type: ignore[explicit-override,override] # pyright: ignore[reportIncompatibleMethodOverride,reportImplicitOverride]
+ raise Exception()
+
+ @jax_autojit
+ def f(x: object) -> object:
+ return x
+
+ inp = C()
+ out = f(inp)
+ assert out is inp
+
+ # Serializable opaque input contains non-serializable object plus array
+ inp = Wrapper((C(), jnp.asarray([1, 2])))
+ out = f(inp)
+ assert isinstance(out, Wrapper)
+ assert out.x[0] is inp.x[0]
+ assert out.x[1] is not inp.x[1]
+ xp_assert_equal(out.x[1], inp.x[1]) # pyright: ignore[reportUnknownArgumentType]
+
+ def test_arraylikes_are_static(self):
+ pytest.importorskip("jax")
+
+ @jax_autojit
+ def f(x: list[int]) -> list[int]:
+ assert isinstance(x, list)
+ assert x == [1, 2]
+ return [3, 4]
+
+ out = f([1, 2])
+ assert isinstance(out, list)
+ assert out == [3, 4]
diff --git a/tests/test_lazy.py b/tests/test_lazy.py
index f40df277..f97062f0 100644
--- a/tests/test_lazy.py
+++ b/tests/test_lazy.py
@@ -15,9 +15,7 @@
from array_api_extra._lib._utils._typing import Array, Device
from array_api_extra.testing import lazy_xp_function
-lazy_xp_function(
- lazy_apply, static_argnames=("func", "shape", "dtype", "as_numpy", "xp")
-)
+lazy_xp_function(lazy_apply)
as_numpy = pytest.mark.parametrize(
"as_numpy",
@@ -280,6 +278,9 @@ def test_lazy_apply_none_shape_broadcast(xp: ModuleType):
Backend.ARRAY_API_STRICT, reason="device->host copy"
),
pytest.mark.skip_xp_backend(Backend.CUPY, reason="device->host copy"),
+ pytest.mark.skip_xp_backend(
+ Backend.TORCH, reason="materialize 'meta' device"
+ ),
pytest.mark.skip_xp_backend(
Backend.TORCH_GPU, reason="device->host copy"
),
@@ -386,7 +387,7 @@ def eager(
)
-lazy_xp_function(check_lazy_apply_kwargs, static_argnames=("expect_cls", "as_numpy"))
+lazy_xp_function(check_lazy_apply_kwargs)
@as_numpy
diff --git a/tests/test_testing.py b/tests/test_testing.py
index fb9ba581..43d6e8ae 100644
--- a/tests/test_testing.py
+++ b/tests/test_testing.py
@@ -6,97 +6,183 @@
import pytest
from array_api_extra._lib._backends import Backend
-from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
+from array_api_extra._lib._testing import (
+ as_numpy_array,
+ xp_assert_close,
+ xp_assert_equal,
+ xp_assert_less,
+)
from array_api_extra._lib._utils._compat import (
array_namespace,
is_dask_namespace,
is_jax_namespace,
)
-from array_api_extra._lib._utils._typing import Array
+from array_api_extra._lib._utils._typing import Array, Device
from array_api_extra.testing import lazy_xp_function
-# mypy: disable-error-code=decorated-any
+# mypy: disable-error-code="decorated-any, explicit-any"
# pyright: reportUnknownParameterType=false,reportMissingParameterType=false
-param_assert_equal_close = pytest.mark.parametrize(
- "func",
- [
- xp_assert_equal,
- pytest.param(
- xp_assert_close,
- marks=pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype"),
- ),
- ],
-)
-
-
-@param_assert_equal_close
-def test_assert_close_equal_basic(xp: ModuleType, func: Callable[..., None]): # type: ignore[explicit-any]
- func(xp.asarray(0), xp.asarray(0))
- func(xp.asarray([1, 2]), xp.asarray([1, 2]))
-
- with pytest.raises(AssertionError, match="shapes do not match"):
- func(xp.asarray([0]), xp.asarray([[0]]))
-
- with pytest.raises(AssertionError, match="dtypes do not match"):
- func(xp.asarray(0, dtype=xp.float32), xp.asarray(0, dtype=xp.float64))
-
- with pytest.raises(AssertionError):
- func(xp.asarray([1, 2]), xp.asarray([1, 3]))
-
- with pytest.raises(AssertionError, match="hello"):
- func(xp.asarray([1, 2]), xp.asarray([1, 3]), err_msg="hello")
-
-
-@pytest.mark.skip_xp_backend(Backend.NUMPY, reason="test other ns vs. numpy")
-@pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="test other ns vs. numpy")
-@pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close])
-def test_assert_close_equal_namespace(xp: ModuleType, func: Callable[..., None]): # type: ignore[explicit-any]
- with pytest.raises(AssertionError, match="namespaces do not match"):
- func(xp.asarray(0), np.asarray(0))
- with pytest.raises(TypeError, match="Unrecognized array input"):
- func(xp.asarray(0), 0)
- with pytest.raises(TypeError, match="list is not a supported array type"):
- func(xp.asarray([0]), [0])
-
-
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
-def test_assert_close_tolerance(xp: ModuleType):
- xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), rtol=0.03)
- with pytest.raises(AssertionError):
- xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), rtol=0.01)
- xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=3)
- with pytest.raises(AssertionError):
- xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=1)
-
-
-@param_assert_equal_close
-@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array")
-@pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="boolean indexing")
-def test_assert_close_equal_none_shape(xp: ModuleType, func: Callable[..., None]): # type: ignore[explicit-any]
- """On Dask and other lazy backends, test that a shape with NaN's or None's
- can be compared to a real shape.
- """
- a = xp.asarray([1, 2])
- a = a[a > 1]
-
- func(a, xp.asarray([2]))
- with pytest.raises(AssertionError):
- func(a, xp.asarray([2, 3]))
- with pytest.raises(AssertionError):
- func(a, xp.asarray(2))
- with pytest.raises(AssertionError):
- func(a, xp.asarray([3]))
-
- # Swap actual and desired
- func(xp.asarray([2]), a)
- with pytest.raises(AssertionError):
- func(xp.asarray([2, 3]), a)
- with pytest.raises(AssertionError):
- func(xp.asarray(2), a)
- with pytest.raises(AssertionError):
- func(xp.asarray([3]), a)
+class TestAsNumPyArray:
+ def test_basic(self, xp: ModuleType):
+ x = xp.asarray([1, 2, 3])
+ y = as_numpy_array(x, xp=xp)
+ xp_assert_equal(y, np.asarray([1, 2, 3])) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+
+ @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device")
+ def test_device(self, xp: ModuleType, device: Device):
+ x = xp.asarray([1, 2, 3], device=device)
+ y = as_numpy_array(x, xp=xp)
+ xp_assert_equal(y, np.asarray([1, 2, 3])) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+
+
+class TestAssertEqualCloseLess:
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close])
+ def test_assert_equal_close_basic(self, xp: ModuleType, func: Callable[..., None]):
+ func(xp.asarray(0), xp.asarray(0))
+ func(xp.asarray([1, 2]), xp.asarray([1, 2]))
+
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(xp.asarray([1, 2]), xp.asarray([2, 1]))
+
+ with pytest.raises(AssertionError, match="hello"):
+ func(xp.asarray([1, 2]), xp.asarray([2, 1]), err_msg="hello")
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ def test_shape_dtype(self, xp: ModuleType, func: Callable[..., None]):
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(xp.asarray([0]), xp.asarray([[0]]))
+
+ with pytest.raises(AssertionError, match="dtypes do not match"):
+ func(xp.asarray(0, dtype=xp.float32), xp.asarray(0, dtype=xp.float64))
+
+ @pytest.mark.skip_xp_backend(Backend.NUMPY, reason="test other ns vs. numpy")
+ @pytest.mark.skip_xp_backend(
+ Backend.NUMPY_READONLY, reason="test other ns vs. numpy"
+ )
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ def test_namespace(self, xp: ModuleType, func: Callable[..., None]):
+ with pytest.raises(AssertionError, match="namespaces do not match"):
+ func(xp.asarray(0), np.asarray(0))
+ with pytest.raises(TypeError, match="Unrecognized array input"):
+ func(xp.asarray(0), 0)
+ with pytest.raises(TypeError, match="list is not a supported array type"):
+ func(xp.asarray([0]), [0])
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ def test_check_shape(self, xp: ModuleType, func: Callable[..., None]):
+ a = xp.asarray([1] if func is xp_assert_less else [2])
+ b = xp.asarray(2)
+ c = xp.asarray(0)
+ d = xp.asarray([2, 2])
+
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(a, b)
+ func(a, b, check_shape=False)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(a, c, check_shape=False)
+ with pytest.raises(AssertionError, match="sizes do not match"):
+ func(a, d, check_shape=False)
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ def test_check_dtype(self, xp: ModuleType, func: Callable[..., None]):
+ a = xp.asarray(1 if func is xp_assert_less else 2)
+ b = xp.asarray(2, dtype=xp.int16)
+ c = xp.asarray(0, dtype=xp.int16)
+
+ with pytest.raises(AssertionError, match="dtypes do not match"):
+ func(a, b)
+ func(a, b, check_dtype=False)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(a, c, check_dtype=False)
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ @pytest.mark.xfail_xp_backend(
+ Backend.SPARSE, reason="sparse [()] returns np.generic"
+ )
+ def test_check_scalar(
+ self, xp: ModuleType, library: Backend, func: Callable[..., None]
+ ):
+ a = xp.asarray(1 if func is xp_assert_less else 2)
+ b = xp.asarray(2)[()] # Note: only makes a difference on NumPy
+ c = xp.asarray(0)
+
+ func(a, b)
+ if library.like(Backend.NUMPY):
+ with pytest.raises(AssertionError, match="array-ness does not match"):
+ func(a, b, check_scalar=True)
+ else:
+ func(a, b, check_scalar=True)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(a, c, check_scalar=True)
+
+ @pytest.mark.parametrize("dtype", ["int64", "float64"])
+ def test_assert_close_tolerance(self, dtype: str, xp: ModuleType):
+ a = xp.asarray([100], dtype=getattr(xp, dtype))
+ b = xp.asarray([102], dtype=getattr(xp, dtype))
+
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ xp_assert_close(a, b)
+
+ xp_assert_close(a, b, rtol=0.03)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ xp_assert_close(a, b, rtol=0.01)
+
+ xp_assert_close(a, b, atol=3)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ xp_assert_close(a, b, atol=1)
+
+ def test_assert_less(self, xp: ModuleType):
+ xp_assert_less(xp.asarray(-1), xp.asarray(0))
+ xp_assert_less(xp.asarray([1, 2]), xp.asarray([2, 3]))
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ xp_assert_less(xp.asarray([1, 1]), xp.asarray([2, 1]))
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="index by sparse array")
+ @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="boolean indexing")
+ def test_none_shape(self, xp: ModuleType, func: Callable[..., None]):
+ """On Dask and other lazy backends, test that a shape with NaN's or None's
+ can be compared to a real shape.
+ """
+ # actual has shape=(None, )
+ a = xp.asarray([1] if func is xp_assert_less else [2])
+ a = a[a > 0]
+
+ func(a, xp.asarray([2]))
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(a, xp.asarray(2))
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(a, xp.asarray([2, 3]))
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(a, xp.asarray([0]))
+
+ # desired has shape=(None, )
+ a = xp.asarray([3] if func is xp_assert_less else [2])
+ a = a[a > 0]
+
+ func(xp.asarray([2]), a)
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(xp.asarray(2), a)
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(xp.asarray([2, 3]), a)
+ with pytest.raises(AssertionError, match="Mismatched elements"):
+ func(xp.asarray([4]), a)
+
+ @pytest.mark.parametrize("func", [xp_assert_equal, xp_assert_close, xp_assert_less])
+ def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None]):
+ a = xp.asarray([1] if func is xp_assert_less else [2], device=device)
+ b = xp.asarray([2], device=device)
+ c = xp.asarray([2, 2], device=device)
+
+ func(a, b)
+ with pytest.raises(AssertionError, match="shapes do not match"):
+ func(a, c)
+ # This is normally performed by np.testing.assert_array_equal etc.
+ # but in case of torch device='meta' we have to do it manually
+ with pytest.raises(AssertionError, match="sizes do not match"):
+ func(a, c, check_shape=False)
def good_lazy(x: Array) -> Array:
@@ -185,32 +271,24 @@ def static_params(x: Array, n: int, flag: bool = False) -> Array:
return x * 3.0
-def static_params1(x: Array, n: int, flag: bool = False) -> Array:
- return static_params(x, n, flag)
-
-
-def static_params2(x: Array, n: int, flag: bool = False) -> Array:
- return static_params(x, n, flag)
-
+lazy_xp_function(static_params)
-def static_params3(x: Array, n: int, flag: bool = False) -> Array:
- return static_params(x, n, flag)
-
-lazy_xp_function(static_params1, static_argnums=(1, 2))
-lazy_xp_function(static_params2, static_argnames=("n", "flag"))
-lazy_xp_function(static_params3, static_argnums=1, static_argnames="flag")
+def test_lazy_xp_function_static_params(xp: ModuleType):
+ x = xp.asarray([1.0, 2.0])
+ xp_assert_equal(static_params(x, 1), xp.asarray([3.0, 6.0]))
+ xp_assert_equal(static_params(x, 1, True), xp.asarray([2.0, 4.0]))
+ xp_assert_equal(static_params(x, 1, False), xp.asarray([3.0, 6.0]))
+ xp_assert_equal(static_params(x, 0, False), xp.asarray([3.0, 6.0]))
+ xp_assert_equal(static_params(x, 1, flag=True), xp.asarray([2.0, 4.0]))
+ xp_assert_equal(static_params(x, n=1, flag=True), xp.asarray([2.0, 4.0]))
-@pytest.mark.parametrize("func", [static_params1, static_params2, static_params3])
-def test_lazy_xp_function_static_params(xp: ModuleType, func: Callable[..., Array]): # type: ignore[explicit-any]
- x = xp.asarray([1.0, 2.0])
- xp_assert_equal(func(x, 1), xp.asarray([3.0, 6.0]))
- xp_assert_equal(func(x, 1, True), xp.asarray([2.0, 4.0]))
- xp_assert_equal(func(x, 1, False), xp.asarray([3.0, 6.0]))
- xp_assert_equal(func(x, 0, False), xp.asarray([3.0, 6.0]))
- xp_assert_equal(func(x, 1, flag=True), xp.asarray([2.0, 4.0]))
- xp_assert_equal(func(x, n=1, flag=True), xp.asarray([2.0, 4.0]))
+def test_lazy_xp_function_deprecated_static_argnames():
+ with pytest.warns(DeprecationWarning, match="static_argnames"):
+ lazy_xp_function(static_params, static_argnames=["flag"]) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+ with pytest.warns(DeprecationWarning, match="static_argnums"):
+ lazy_xp_function(static_params, static_argnums=[1]) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
try:
@@ -273,6 +351,66 @@ def test_lazy_xp_function_eagerly_raises(da: ModuleType):
_ = dask_raises(x)
+class Wrapper:
+ """Trivial opaque wrapper. Must be pickleable."""
+
+ x: Array
+
+ def __init__(self, x: Array):
+ self.x = x
+
+
+def check_opaque_wrapper(w: Wrapper, xp: ModuleType) -> Wrapper:
+ assert isinstance(w, Wrapper)
+ assert array_namespace(w.x) == xp
+ return Wrapper(w.x + 1)
+
+
+lazy_xp_function(check_opaque_wrapper)
+
+
+def test_lazy_xp_function_opaque_wrappers(xp: ModuleType):
+ """
+ Test that function input and output can be wrapped into arbitrary
+ serializable Python objects, even if jax.jit does not support them.
+ """
+ x = xp.asarray([1, 2])
+ xp2 = array_namespace(x) # Revert NUMPY_READONLY to array_api_compat.numpy
+ res = check_opaque_wrapper(Wrapper(x), xp2)
+ xp_assert_equal(res.x, xp.asarray([2, 3]))
+
+
+def test_lazy_xp_function_opaque_wrappers_eagerly_raise(da: ModuleType):
+ """
+ Like `test_lazy_xp_function_eagerly_raises`, but the returned object is
+ wrapped in an opaque wrapper.
+ """
+ x = da.arange(3)
+ with pytest.raises(ValueError, match="Hello world"):
+ _ = Wrapper(dask_raises(x))
+
+
+def check_recursive(x: list[object]) -> list[object]:
+ assert isinstance(x, list)
+ assert x[1] is x
+ y: list[object] = [cast(Array, x[0]) + 1]
+ y.append(y)
+ return y
+
+
+lazy_xp_function(check_recursive)
+
+
+def test_lazy_xp_function_recursive(xp: ModuleType):
+ """Test that inputs and outputs can be recursive data structures."""
+ x: list[object] = [xp.asarray([1, 2])]
+ x.append(x)
+ y = check_recursive(x)
+ assert isinstance(y, list)
+ xp_assert_equal(cast(Array, y[0]), xp.asarray([2, 3]))
+ assert y[1] is y
+
+
wrapped = ModuleType("wrapped")
naked = ModuleType("naked")
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy