|
| 1 | +""" |
| 2 | +Public testing utilities. |
| 3 | +
|
| 4 | +See also _lib._testing for additional private testing utilities. |
| 5 | +""" |
| 6 | + |
| 7 | +# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972 |
| 8 | +from __future__ import annotations |
| 9 | + |
| 10 | +from collections.abc import Callable, Iterable, Sequence |
| 11 | +from functools import wraps |
| 12 | +from types import ModuleType |
| 13 | +from typing import TYPE_CHECKING, Any, TypeVar, cast |
| 14 | + |
| 15 | +import pytest |
| 16 | + |
| 17 | +from array_api_extra._lib._utils._compat import is_dask_namespace, is_jax_namespace |
| 18 | + |
| 19 | +__all__ = ["lazy_xp_function", "patch_lazy_xp_functions"] |
| 20 | + |
| 21 | +if TYPE_CHECKING: |
| 22 | + # TODO move outside TYPE_CHECKING |
| 23 | + # depends on scikit-learn abandoning Python 3.9 |
| 24 | + # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972 |
| 25 | + from typing import ParamSpec |
| 26 | + |
| 27 | + P = ParamSpec("P") |
| 28 | +else: |
| 29 | + # Sphinx hacks |
| 30 | + class P: # pylint: disable=missing-class-docstring |
| 31 | + args: tuple |
| 32 | + kwargs: dict |
| 33 | + |
| 34 | + |
| 35 | +T = TypeVar("T") |
| 36 | + |
| 37 | + |
| 38 | +def lazy_xp_function( # type: ignore[no-any-explicit] |
| 39 | + func: Callable[..., Any], |
| 40 | + *, |
| 41 | + dask_disable_compute: bool = True, |
| 42 | + jax_jit: bool = True, |
| 43 | + static_argnums: int | Sequence[int] | None = None, |
| 44 | + static_argnames: str | Iterable[str] | None = None, |
| 45 | +) -> None: # numpydoc ignore=GL07 |
| 46 | + """ |
| 47 | + Tag a function to be tested for lazy backends. |
| 48 | +
|
| 49 | + Tag a function, which must be imported in the test module globals, so that when any |
| 50 | + tests defined in the same module are executed with `xp=jax.numpy` the function is |
| 51 | + replaced with a jitted version of itself, and when it is executed with |
| 52 | + `xp=dask.array` the function will raise if it attempts to materialize the graph. |
| 53 | +
|
| 54 | + This will be later expanded to provide test coverage for other lazy backends. |
| 55 | +
|
| 56 | + Parameters |
| 57 | + ---------- |
| 58 | + func : callable |
| 59 | + Function to be tested. |
| 60 | + dask_disable_compute : bool, optional |
| 61 | + Set to True to raise an error if `func` attempts to call `dask.compute()` or |
| 62 | + `dask.persist()`. This is typically inadvertently triggered by `bool()`, |
| 63 | + `float()`, and `np.asarray()`. Set to False to allow these calls, knowing that |
| 64 | + they are going to be extremely detrimental for performance. |
| 65 | + jax_jit : bool, optional |
| 66 | + Set to True to replace `func` with `jax.jit(func)` when calling the |
| 67 | + `patch_lazy_xp_functions` test helper with `xp=jax.numpy`. |
| 68 | + Set to False if `func` is only compatible with eager (non-jitted) JAX. |
| 69 | + Default: True. |
| 70 | + static_argnums : int | Sequence[int], optional |
| 71 | + Passed to jax.jit. |
| 72 | + Positional arguments to treat as static (trace- and compile-time constant). |
| 73 | + Default: infer from static_argnames using `inspect.signature(func)`. |
| 74 | + static_argnames : str | Iterable[str], optional |
| 75 | + Passed to jax.jit. |
| 76 | + Named arguments to treat as static (compile-time constant). |
| 77 | + Default: infer from static_argnums using `inspect.signature(func)`. |
| 78 | +
|
| 79 | + See Also |
| 80 | + -------- |
| 81 | + patch_lazy_xp_functions |
| 82 | + jax.jit |
| 83 | +
|
| 84 | + Examples |
| 85 | + -------- |
| 86 | + In `test_mymodule.py`:: |
| 87 | +
|
| 88 | + from array_api_extra.testing import lazy_xp_function |
| 89 | + from mymodule import myfunc |
| 90 | +
|
| 91 | + lazy_xp_function(myfunc) |
| 92 | +
|
| 93 | + def test_myfunc(xp): |
| 94 | + a = xp.asarray([1, 2]) |
| 95 | + # When xp=jax.numpy, this is the same as `b = jax.jit(myfunc)(a)` |
| 96 | + # When xp=dask.array, crash on compute() or persist() |
| 97 | + b = myfunc(a) |
| 98 | +
|
| 99 | + Notes |
| 100 | + ----- |
| 101 | + A test function can circumvent this monkey-patching system by calling `func` an |
| 102 | + attribute of the original module. You need to sanitize your code to |
| 103 | + make sure this does not happen. |
| 104 | +
|
| 105 | + Example:: |
| 106 | +
|
| 107 | + import mymodule |
| 108 | + from mymodule import myfunc |
| 109 | +
|
| 110 | + lazy_xp_function(myfunc) |
| 111 | +
|
| 112 | + def test_myfunc(xp): |
| 113 | + a = xp.asarray([1, 2]) |
| 114 | + b = myfunc(a) # This is jitted when xp=jax.numpy |
| 115 | + c = mymodule.myfunc(a) # This is not |
| 116 | + """ |
| 117 | + func.dask_disable_compute = dask_disable_compute # type: ignore[attr-defined] # pyright: ignore[reportFunctionMemberAccess] |
| 118 | + if jax_jit: |
| 119 | + func.lazy_jax_jit_kwargs = { # type: ignore[attr-defined] # pyright: ignore[reportFunctionMemberAccess] |
| 120 | + "static_argnums": static_argnums, |
| 121 | + "static_argnames": static_argnames, |
| 122 | + } |
| 123 | + |
| 124 | + |
| 125 | +def patch_lazy_xp_functions( |
| 126 | + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, *, xp: ModuleType |
| 127 | +) -> None: |
| 128 | + """ |
| 129 | + Test lazy execution. |
| 130 | +
|
| 131 | + If `xp==jax.numpy`, search for all functions which have been tagged by |
| 132 | + `lazy_xp_function` in the globals of the module that defines the current test |
| 133 | + and wrap them with `jax.jit`. Unwrap them at the end of the test. |
| 134 | +
|
| 135 | + If `xp==dask.array`, wrap the functions with a decorator that disables `compute()` |
| 136 | + and `persist()`. |
| 137 | +
|
| 138 | + This function should be called by your library's `xp` fixture that runs tests on |
| 139 | + multiple backends:: |
| 140 | +
|
| 141 | + @pytest.fixture(params=[numpy, array_api_strict, jax.numpy, dask.array]) |
| 142 | + def xp(request, monkeypatch): |
| 143 | + patch_lazy_xp_functions(request, monkeypatch, xp=request.param) |
| 144 | + return request.param |
| 145 | +
|
| 146 | + Parameters |
| 147 | + ---------- |
| 148 | + request : pytest.FixtureRequest |
| 149 | + Pytest fixture, as acquired by the test itself or by one of its fixtures. |
| 150 | + monkeypatch : pytest.MonkeyPatch |
| 151 | + Pytest fixture, as acquired by the test itself or by one of its fixtures. |
| 152 | + xp : module |
| 153 | + Array namespace to be tested. |
| 154 | +
|
| 155 | + See Also |
| 156 | + -------- |
| 157 | + lazy_xp_function |
| 158 | + pytest.FixtureRequest |
| 159 | + """ |
| 160 | + globals_ = cast(dict[str, Any], request.module.__dict__) # type: ignore[no-any-explicit] |
| 161 | + |
| 162 | + if is_dask_namespace(xp): |
| 163 | + for name, func in globals_.items(): |
| 164 | + if getattr(func, "dask_disable_compute", False): |
| 165 | + wrapped = _dask_disable_compute(func) |
| 166 | + monkeypatch.setitem(globals_, name, wrapped) |
| 167 | + |
| 168 | + elif is_jax_namespace(xp): |
| 169 | + import jax |
| 170 | + |
| 171 | + for name, func in globals_.items(): |
| 172 | + kwargs = cast( # type: ignore[no-any-explicit] |
| 173 | + "dict[str, Any] | None", getattr(func, "lazy_jax_jit_kwargs", None) |
| 174 | + ) |
| 175 | + |
| 176 | + # suppress unused-ignore to run mypy in -e lint as well as -e dev |
| 177 | + if kwargs is not None: # type: ignore[no-untyped-call,unused-ignore] |
| 178 | + wrapped = jax.jit(func, **kwargs) # type: ignore[no-untyped-call,unused-ignore] |
| 179 | + monkeypatch.setitem(globals_, name, wrapped) # pyright: ignore[reportUnknownArgumentType] |
| 180 | + |
| 181 | + |
| 182 | +def _dask_disable_compute( |
| 183 | + func: Callable[P, T], |
| 184 | +) -> Callable[P, T]: # numpydoc ignore=PR01,RT01 |
| 185 | + """ |
| 186 | + Wrap a function to raise if it attempts to call dask.compute or dask.persist. |
| 187 | + """ |
| 188 | + import dask.config |
| 189 | + |
| 190 | + def get(*args: object, **kwargs: object) -> object: # noqa: ARG001 # numpydoc ignore=PR01 |
| 191 | + """Dask scheduler which will always raise when invoked.""" |
| 192 | + msg = "Called `dask.compute()` or `dask.persist()`" |
| 193 | + raise AssertionError(msg) |
| 194 | + |
| 195 | + @wraps(func) |
| 196 | + def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: # numpydoc ignore=GL08 |
| 197 | + with dask.config.set({"scheduler": get}): |
| 198 | + return func(*args, **kwargs) |
| 199 | + |
| 200 | + return wrapper |
0 commit comments