BUG: lazy_xp_function crashes with ufuncs

Author: crusaderky

src/array_api_extra/testing.py

@@ -7,7 +7,8 @@
 # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Sequence
+import sys
+from collections.abc import Callable, Iterable, Iterator, Sequence
 from functools import wraps
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, TypeVar, cast
@@ -42,6 +43,8 @@ def override(func: Callable[P, T]) -> Callable[P, T]:
 
 T = TypeVar("T")
 
+_ufuncs_tags: dict[object, dict[str, Any]] = {}  # type: ignore[no-any-explicit]
+
 
 def lazy_xp_function(  # type: ignore[no-any-explicit]
     func: Callable[..., Any],
@@ -132,12 +135,30 @@ 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
     """
-    func.allow_dask_compute = allow_dask_compute  # type: ignore[attr-defined]  # pyright: ignore[reportFunctionMemberAccess]
-    if jax_jit:
-        func.lazy_jax_jit_kwargs = {  # type: ignore[attr-defined]  # pyright: ignore[reportFunctionMemberAccess]
-            "static_argnums": static_argnums,
-            "static_argnames": static_argnames,
-        }
+    tags = {
+        "allow_dask_compute": allow_dask_compute,
+        "jax_jit": jax_jit,
+        "static_argnums": static_argnums,
+        "static_argnames": static_argnames,
+    }
+    if _is_numpy_ufunc(func):
+        _ufuncs_tags[func] = tags  # Can't assign attributes
+    else:
+        func._lazy_xp_function = tags  # type: ignore[attr-defined]  # pylint: disable=protected-access  # pyright: ignore[reportFunctionMemberAccess]
+
+
+def _is_numpy_ufunc(f: Callable[..., Any]) -> bool:  # type: ignore[no-any-explicit]
+    """Check if a function is a numpy ufunc
+    native from numpy, CPython extension, or Cythonized).
+    """
+    if "numpy" not in sys.modules:
+        # FIXME: is it possible to import a CPython or Cython module
+        # that defines a ufunc without importing numpy first?
+        return False
+
+    import numpy as np
+
+    return isinstance(f, np.ufunc)
 
 
 def patch_lazy_xp_functions(
@@ -179,24 +200,35 @@ def xp(request, monkeypatch):
     """
     globals_ = cast("dict[str, Any]", request.module.__dict__)  # type: ignore[no-any-explicit]
 
-    if is_dask_namespace(xp):
+    def iter_tagged() -> Iterator[tuple[str, Callable[..., Any], dict[str, Any]]]:  # type: ignore[no-any-explicit]
         for name, func in globals_.items():
-            n = getattr(func, "allow_dask_compute", None)
-            if n is not None:
-                assert isinstance(n, int)
-                wrapped = _allow_dask_compute(func, n)
-                monkeypatch.setitem(globals_, name, wrapped)
+            if _is_numpy_ufunc(func):
+                tags = _ufuncs_tags.get(func)
+            else:
+                tags = getattr(func, "_lazy_xp_function", None)  # type: ignore[assignment]
+            if tags:
+                yield name, func, tags
+
+    if is_dask_namespace(xp):
+        for name, func, tags in iter_tagged():
+            n = tags["allow_dask_compute"]
+            wrapped = _allow_dask_compute(func, n)
+            monkeypatch.setitem(globals_, name, wrapped)
 
     elif is_jax_namespace(xp):
         import jax
 
-        for name, func in globals_.items():
-            kwargs = cast(  # type: ignore[no-any-explicit]
-                "dict[str, Any] | None", getattr(func, "lazy_jax_jit_kwargs", None)
-            )
-            if kwargs is not None:
+        for 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(Callable[..., Any], jax.jit(func, **kwargs))  # type: ignore[no-any-explicit,no-untyped-call,unused-ignore]
+                wrapped = cast(  # type: ignore[no-any-explicit]
+                    Callable[..., Any],
+                    jax.jit(
+                        func,
+                        static_argnums=tags["static_argnums"],
+                        static_argnames=tags["static_argnames"],
+                    ),
+                )
                 monkeypatch.setitem(globals_, name, wrapped)
 
 

tests/test_testing.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+from numpy import min as numpy_min
 
 from array_api_extra._lib import Backend
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
@@ -202,3 +203,23 @@ def test_lazy_xp_function_static_params(xp: ModuleType, func: Callable[..., Arra
     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]))
+
+
+lazy_xp_function(numpy_min, static_argnames="axis")
+
+
+def test_lazy_xp_function_ufunc(xp: ModuleType, library: Backend):
+    x = xp.asarray([[1, 4], [3, 2]])
+    if library in (Backend.ARRAY_API_STRICT, Backend.TORCH, Backend.JAX):
+        # array-api-strict, torch and jax don't define __array_ufunc__
+        # numpy ufuncs can't auto-convert to numpy from torch
+        # array-api-strict arrays are auto-converted to numpy
+        # eager jax arrays are auto-converted to numpy in eager jax
+        # and fail in jax.jit (which lazy_xp_function tests here)
+        with pytest.raises((TypeError, AssertionError)):
+            xp_assert_equal(numpy_min(x, axis=0), xp.asarray([1, 2]))
+    else:
+        # cupy, dask and sparse define __array_ufunc__ and dispatch accordingly
+        # note that when sparse reduces to scalar it returns a np.generic, which
+        # would make xp_assert_equal fail.
+        xp_assert_equal(numpy_min(x, axis=0), xp.asarray([1, 2]))