Merge branch 'main' into dask-new

[skip cirrus] [skip circle]

Author: lucascolley

doc/source/dev/api-dev/array_api.rst

@@ -331,6 +331,32 @@ Note that there is a GitHub Actions workflow which tests with array-api-strict,
 PyTorch, and JAX on CPU.
 
 
+Testing the JAX JIT compiler
+----------------------------
+The `JAX JIT compiler <https://jax.readthedocs.io/en/latest/jit-compilation.html>`_
+introduces special restrictions to all code wrapped by `@jax.jit`, which are not
+present when running JAX in eager mode. Notably, boolean masks in `__getitem__`
+and `.at` aren't supported, and you can't materialize the arrays by applying
+`bool()`, `float()`, `np.asarray()` etc. to them.
+
+To properly test scipy with JAX, you need to wrap the tested scipy functions
+with `@jax.jit` before they are called by the unit tests.
+To achieve this, you should tag them as follows in your test module::
+
+  from scipy._lib._lazy_testing import lazy_xp_function
+  from scipy.mymodule import toto
+
+  lazy_xp_function(toto)
+
+  def test_toto(xp):
+      a = xp.asarray([1, 2, 3])
+      b = xp.asarray([0, 2, 5])
+      # When xp==jax.numpy, toto is wrapped with @jax.jit
+      xp_assert_close(toto(a, b), a)
+
+See full documentation in `scipy/_lib/_lazy_testing.py`.
+
+
 Additional information
 ----------------------
 

mypy.ini

@@ -88,6 +88,10 @@ ignore_missing_imports = True
 [mypy-array_api_strict]
 ignore_missing_imports = True
 
+[mypy-jax]
+# Typed, but cumbersome to install in CI (depends on scipy)
+ignore_missing_imports = True
+
 [mypy-sphinx.*]
 ignore_missing_imports = True
 
@@ -751,3 +755,6 @@ ignore_errors = True
 
 [mypy-scipy._lib.array_api_compat.*]
 ignore_errors = True
+
+[mypy-scipy._lib.array_api_extra.*]
+ignore_errors = True

scipy/_lib/_lazy_testing.py

@@ -0,0 +1,111 @@
+from collections.abc import Callable, Iterable, Sequence
+from types import ModuleType
+import pytest
+from scipy._lib._array_api import is_jax
+
+
+def lazy_xp_function(
+    func: Callable,
+    *,
+    jax_jit: bool = True,
+    static_argnums: int | Sequence[int] | None=None,
+    static_argnames: str | Iterable[str] | None=None,
+) -> None:
+    """Tag a function, which must be imported in the test module globals,
+    so that when any tests defined in the same module are executed with
+    xp=jax.numpy the function is replaced with a jitted version of itself.
+
+    This will be later expanded to provide test coverage for other lazy backends,
+    e.g. Dask.
+
+    Example::
+
+      # test_mymodule.py:
+      from scipy._lib._lazy_testing import lazy_xp_function
+      from scipy.mymodule import myfunc
+
+      lazy_xp_function(myfunc)
+
+      def test_myfunc(xp):
+          a = xp.asarray([1, 2])
+          # When xp=jax.numpy, this is the same as
+          # b = jax.jit(myfunc)(a)
+          b = myfunc(a)
+
+    Parameters
+    ----------
+    func : callable
+        Function to be tested
+    jax_jit : bool, optional
+        Set to True to replace `func` with `jax.jit(func)` when calling the
+        `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 (trace- and 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)`.
+
+    Notes
+    -----
+    A test function can circumvent this monkey-patching system by calling `func` an
+    attribute of the original module. You need to sanitize your code to
+    make sure this does not happen.
+
+    Example::
+
+      import mymodule
+      from mymodule import myfunc
+
+      lazy_xp_function(myfunc)
+
+      def test_myfunc(xp):
+          a = xp.asarray([1, 2])
+          b = myfunc(a)  # This is jitted when xp=jax.numpy
+          c = mymodule.myfunc(a)  # This is not
+
+    See Also
+    --------
+    patch_lazy_xp_functions
+    jax.jit: https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html
+    """
+    if jax_jit:
+        func._lazy_jax_jit_kwargs = {  # type: ignore[attr-defined]
+            "static_argnums": static_argnums,
+            "static_argnames": static_argnames,
+        }
+
+
+def patch_lazy_xp_functions(
+    xp: ModuleType, request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """If xp==jax.numpy, search for all functions which have been tagged by
+    `lazy_xp_function` in the globals of the module that defines the current test
+    and wrap them with `jax.jit`. Unwrap them at the end of the test.
+
+    Parameters
+    ----------
+    xp: module
+        Array namespace to be tested
+    request: pytest.FixtureRequest
+        Pytest fixture, as acquired by the test itself or by one of its fixtures.
+    monkeypatch: pytest.MonkeyPatch
+        Pytest fixture, as acquired by the test itself or by one of its fixtures.
+
+    See Also
+    --------
+    lazy_xp_function
+    https://docs.pytest.org/en/stable/reference/reference.html#std-fixture-request
+    """
+    if is_jax(xp):
+        import jax
+
+        globals_ = request.module.__dict__
+        for name, func in globals_.items():
+            kwargs = getattr(func, "_lazy_jax_jit_kwargs", None)
+            if kwargs is not None:
+                monkeypatch.setitem(globals_, name, jax.jit(func, **kwargs))

scipy/_lib/array_api_extra

@@ -122,6 +122,7 @@ python_sources = [
   '_elementwise_iterative_method.py',
   '_finite_differences.py',
   '_gcutils.py',
+  '_lazy_testing.py',
   '_pep440.py',
   '_testutils.py',
   '_threadsafety.py',
@@ -217,21 +218,32 @@ py3.install_sources(
 # `array_api_extra` install to simplify import path;
 # should be updated whenever new files are added to `array_api_extra`
 
+py3.install_sources(
+  [
+    'array_api_extra/src/array_api_extra/_lib/_utils/__init__.py',
+    'array_api_extra/src/array_api_extra/_lib/_utils/_compat.py',
+    'array_api_extra/src/array_api_extra/_lib/_utils/_compat.pyi',    
+    'array_api_extra/src/array_api_extra/_lib/_utils/_helpers.py',
+    'array_api_extra/src/array_api_extra/_lib/_utils/_typing.py',
+
+  ],
+  subdir: 'scipy/_lib/array_api_extra/_lib/_utils',
+)
+
 py3.install_sources(
   [
     'array_api_extra/src/array_api_extra/_lib/__init__.py',
-    'array_api_extra/src/array_api_extra/_lib/_compat.py',
-    'array_api_extra/src/array_api_extra/_lib/_compat.pyi',    
-    'array_api_extra/src/array_api_extra/_lib/_utils.py',
-    'array_api_extra/src/array_api_extra/_lib/_typing.py',
+    'array_api_extra/src/array_api_extra/_lib/_backends.py',
+    'array_api_extra/src/array_api_extra/_lib/_funcs.py',    
+    'array_api_extra/src/array_api_extra/_lib/_testing.py',
   ],
   subdir: 'scipy/_lib/array_api_extra/_lib',
 )
 
 py3.install_sources(
   [
     'array_api_extra/src/array_api_extra/__init__.py',
-    'array_api_extra/src/array_api_extra/_funcs.py',
+    'array_api_extra/src/array_api_extra/_delegation.py',
   ],
   subdir: 'scipy/_lib/array_api_extra',
 )

scipy/_lib/meson.build

@@ -11,6 +11,7 @@ python_sources = [
   'test_deprecation.py',
   'test_doccer.py',
   'test_import_cycles.py',
+  'test_lazy_testing.py',
   'test_public_api.py',
   'test_scipy_version.py',
   'test_tmpdirs.py',

scipy/_lib/tests/meson.build

@@ -14,6 +14,7 @@
 
 from scipy._lib._array_api import (xp_assert_equal, xp_assert_close, is_numpy,
                                    is_array_api_strict)
+from scipy._lib._lazy_testing import lazy_xp_function
 from scipy._lib._util import (_aligned_zeros, check_random_state, MapWrapper,
                               getfullargspec_no_self, FullArgSpec,
                               rng_integers, _validate_int, _rename_parameter,
@@ -23,6 +24,11 @@
 
 skip_xp_backends = pytest.mark.skip_xp_backends
 
+lazy_xp_function(_contains_nan, static_argnames=("nan_policy", "xp_omit_okay", "xp"))
+# FIXME @jax.jit fails: complex bool mask
+lazy_xp_function(_lazywhere, jax_jit=False, static_argnames=("f", "f2"))
+
+
 @pytest.mark.slow
 def test__aligned_zeros():
     niter = 10
@@ -344,6 +350,7 @@ def test_contains_nan_with_strings(self):
         data4 = np.array([["1", 2], [3, np.nan]], dtype='object')
         assert _contains_nan(data4)
 
+    @pytest.mark.skip_xp_backends("jax.numpy", reason="lazy backends tested separately")
     @pytest.mark.parametrize("nan_policy", ['propagate', 'omit', 'raise'])
     def test_array_api(self, xp, nan_policy):
         rng = np.random.default_rng(932347235892482)
@@ -359,9 +366,40 @@ def test_array_api(self, xp, nan_policy):
         elif nan_policy == 'omit' and not is_numpy(xp):
             with pytest.raises(ValueError, match="nan_policy='omit' is incompatible"):
                 _contains_nan(x, nan_policy)
+            assert _contains_nan(x, nan_policy, xp_omit_okay=True)
         elif nan_policy == 'propagate':
             assert _contains_nan(x, nan_policy)
 
+    @pytest.mark.skip_xp_backends("numpy", reason="lazy backends only")
+    @pytest.mark.skip_xp_backends("cupy", reason="lazy backends only")
+    @pytest.mark.skip_xp_backends("array_api_strict", reason="lazy backends only")
+    @pytest.mark.skip_xp_backends("torch", reason="lazy backends only")
+    def test_array_api_lazy(self, xp):
+        rng = np.random.default_rng(932347235892482)
+        x0 = rng.random(size=(2, 3, 4))
+        x = xp.asarray(x0)
+
+        xp_assert_equal(_contains_nan(x), xp.asarray(False))
+        xp_assert_equal(_contains_nan(x, "propagate"), xp.asarray(False))
+        xp_assert_equal(_contains_nan(x, "omit", xp_omit_okay=True), xp.asarray(False))
+        # Lazy arrays don't support "omit" and "raise" policies
+        # TODO test that we're emitting a user-friendly error message.
+        #      Blocked by https://github.com/data-apis/array-api-compat/pull/228
+        with pytest.raises(TypeError):
+            _contains_nan(x, "omit")
+        with pytest.raises(TypeError):
+            _contains_nan(x, "raise")
+
+        x = xpx.at(x)[1, 2, 1].set(np.nan)
+
+        xp_assert_equal(_contains_nan(x), xp.asarray(True))
+        xp_assert_equal(_contains_nan(x, "propagate"), xp.asarray(True))
+        xp_assert_equal(_contains_nan(x, "omit", xp_omit_okay=True), xp.asarray(True))
+        with pytest.raises(TypeError):
+            _contains_nan(x, "omit")
+        with pytest.raises(TypeError):
+            _contains_nan(x, "raise")
+
 
 def test__rng_html_rewrite():
     def mock_str():

scipy/_lib/tests/test__util.py

@@ -7,6 +7,11 @@
 )
 from scipy._lib import array_api_extra as xpx
 from scipy._lib._array_api_no_0d import xp_assert_equal as xp_assert_equal_no_0d
+from scipy._lib._lazy_testing import lazy_xp_function
+
+lazy_xp_function(_asarray, static_argnames=(
+                 "dtype", "order", "copy", "xp", "check_finite", "subok"))
+lazy_xp_function(xp_copy, static_argnames=("xp", ))
 
 skip_xp_backends = pytest.mark.skip_xp_backends
 

scipy/_lib/tests/test_array_api.py

@@ -0,0 +1,75 @@
+import pytest
+from scipy._lib._array_api import array_namespace, is_jax, xp_assert_equal
+from scipy._lib._lazy_testing import lazy_xp_function
+
+
+def jittable(x):
+    """A jittable function"""
+    return x * 2.0
+
+
+def non_jittable(x):
+    """This function materializes the input array, so it will fail
+    when wrapped in jax.jit
+    """
+    xp = array_namespace(x)
+    if xp.any(x < 0.0):
+        raise ValueError("Negative values not allowed")
+    return x
+
+
+def non_jittable2(x):
+    return non_jittable(x)
+
+
+def static_params(x, n, flag=False):
+    """Function with static parameters that must not be jitted"""
+    if flag and n > 0:  # This fails if n or flag are jitted arrays
+        return x * 2.0
+    else:
+        return x * 3.0
+
+
+def static_params1(x, n, flag=False):
+    return static_params(x, n, flag)
+
+
+def static_params2(x, n, flag=False):
+    return static_params(x, n, flag)
+
+
+def static_params3(x, n, flag=False):
+    return static_params(x, n, flag)
+
+
+lazy_xp_function(jittable)
+lazy_xp_function(non_jittable2)
+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(xp):
+    x = xp.asarray([1.0, 2.0])
+
+    xp_assert_equal(jittable(x), xp.asarray([2.0, 4.0]))
+
+    xp_assert_equal(non_jittable(x), xp.asarray([1.0, 2.0]))  # Not jitted
+    if is_jax(xp):
+        with pytest.raises(
+            TypeError, match="Attempted boolean conversion of traced array"
+        ):
+            non_jittable2(x)  # Jitted
+    else:
+        xp_assert_equal(non_jittable2(x), xp.asarray([1.0, 2.0]))
+
+
+@pytest.mark.parametrize("func", [static_params1, static_params2, static_params3])
+def test_lazy_xp_function_static_params(xp, func):
+    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]))