array-api-extra 0.10.0__tar.gz → 0.10.1__tar.gz

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/.all-contributorsrc

@@ -308,6 +308,17 @@
         "ideas",
         "test"
       ]
+    },
+    {
+      "login": "steppi",
+      "name": "Albert Steppi",
+      "avatar_url": "https://avatars.githubusercontent.com/u/1953382?v=4",
+      "profile": "http://steppi.github.io",
+      "contributions": [
+        "code",
+        "ideas",
+        "test"
+      ]
     }
   ]
 }

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/CONTRIBUTORS.md

@@ -44,6 +44,9 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/adriagarp"><img src="https://avatars.githubusercontent.com/u/96059447?v=4?s=100" width="100px;" alt="Adrián García Pitarch"/><br /><sub><b>Adrián García Pitarch</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/cakedev0"><img src="https://avatars.githubusercontent.com/u/25986961?v=4?s=100" width="100px;" alt="Arthur Lacote"/><br /><sub><b>Arthur Lacote</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Documentation">📖</a> <a href="#ideas-cakedev0" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Tests">⚠️</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="http://steppi.github.io"><img src="https://avatars.githubusercontent.com/u/1953382?v=4?s=100" width="100px;" alt="Albert Steppi"/><br /><sub><b>Albert Steppi</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Code">💻</a> <a href="#ideas-steppi" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Tests">⚠️</a></td>
+    </tr>
   </tbody>
 </table>
 

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: array-api-extra
-Version: 0.10.0
+Version: 0.10.1
 Summary: Extra array functions built on top of the array API standard.
 Author-Email: Lucas Colley <lucas.colley8@gmail.com>
 License-Expression: MIT
@@ -124,6 +124,9 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/adriagarp"><img src="https://avatars.githubusercontent.com/u/96059447?v=4?s=100" width="100px;" alt="Adrián García Pitarch"/><br /><sub><b>Adrián García Pitarch</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/cakedev0"><img src="https://avatars.githubusercontent.com/u/25986961?v=4?s=100" width="100px;" alt="Arthur Lacote"/><br /><sub><b>Arthur Lacote</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Documentation">📖</a> <a href="#ideas-cakedev0" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Tests">⚠️</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="http://steppi.github.io"><img src="https://avatars.githubusercontent.com/u/1953382?v=4?s=100" width="100px;" alt="Albert Steppi"/><br /><sub><b>Albert Steppi</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Code">💻</a> <a href="#ideas-steppi" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Tests">⚠️</a></td>
+    </tr>
   </tbody>
 </table>
 

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/README.md

@@ -99,6 +99,9 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/adriagarp"><img src="https://avatars.githubusercontent.com/u/96059447?v=4?s=100" width="100px;" alt="Adrián García Pitarch"/><br /><sub><b>Adrián García Pitarch</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=adriagarp" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/cakedev0"><img src="https://avatars.githubusercontent.com/u/25986961?v=4?s=100" width="100px;" alt="Arthur Lacote"/><br /><sub><b>Arthur Lacote</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Documentation">📖</a> <a href="#ideas-cakedev0" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=cakedev0" title="Tests">⚠️</a></td>
     </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="http://steppi.github.io"><img src="https://avatars.githubusercontent.com/u/1953382?v=4?s=100" width="100px;" alt="Albert Steppi"/><br /><sub><b>Albert Steppi</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Code">💻</a> <a href="#ideas-steppi" title="Ideas, Planning, & Feedback">🤔</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=steppi" title="Tests">⚠️</a></td>
+    </tr>
   </tbody>
 </table>
 

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/meson.build

@@ -1,6 +1,6 @@
 project(
   'array-api-extra',
-  version: '0.10.0'
+  version: '0.10.1'
 )
 
 py = import('python').find_installation()

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "mesonpy"
 
 [project]
 name = "array-api-extra"
-version = "0.10.0"
+version = "0.10.1"
 authors = [
   { name = "Lucas Colley", email = "lucas.colley8@gmail.com" },
   # { name = "Open Source Contributors" },  # https://github.com/pypi/warehouse/issues/14813
@@ -93,7 +93,7 @@ array-api-extra.path = "."
 typing-extensions = ">=4.15.0"
 pylint = ">=4.0.4"
 mypy = ">=1.19.1"
-basedpyright = ">=1.37.4"
+basedpyright = ">=1.38.0"
 numpydoc = ">=1.10.0,<2"
 # import dependencies for mypy:
 array-api-strict = ">=2.4.1,<2.5"
@@ -101,9 +101,9 @@ numpy = ">=2.1.3"
 hypothesis = ">=6.151.2"
 dask-core = ">=2026.1.2" # No distributed, tornado, etc.
 dprint = ">=0.50.0,<0.51"
-lefthook = ">=2.1.0,<3"
-ruff = ">=0.15.0,<0.16"
-typos = ">=1.43.3,<2"
+lefthook = ">=2.1.1,<3"
+ruff = ">=0.15.1,<0.16"
+typos = ">=1.43.4,<2"
 actionlint = ">=1.7.10,<2"
 blacken-docs = ">=1.20.0,<2"
 pytest = ">=9.0.2,<10"
@@ -134,6 +134,7 @@ pytest-cov = ">=7.0.0"
 hypothesis = ">=6.151.2"
 array-api-strict = ">=2.4.1,<2.5"
 numpy = ">=1.22.0"
+scipy = ">=1.15.2,<2"
 
 [tool.pixi.feature.tests.tasks]
 tests = { cmd = "pytest -v", description = "Run tests" }

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/src/array_api_extra/__init__.py

@@ -27,7 +27,7 @@ from ._lib._funcs import (
 )
 from ._lib._lazy import lazy_apply
 
-__version__ = "0.10.0"
+__version__ = "0.10.1"
 
 # pylint: disable=duplicate-code
 __all__ = [

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/src/array_api_extra/testing.py

@@ -10,8 +10,9 @@ import contextlib
 import enum
 import warnings
 from collections.abc import Callable, Generator, Iterator, Sequence
-from functools import wraps
-from types import ModuleType
+from functools import update_wrapper, wraps
+from inspect import getattr_static
+from types import FunctionType, ModuleType
 from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
 
 from ._lib._utils._compat import is_dask_namespace, is_jax_namespace
@@ -48,8 +49,21 @@ class Deprecated(enum.Enum):
 DEPRECATED = Deprecated.DEPRECATED
 
 
+def _clone_function(f: Callable[..., Any]) -> Callable[..., Any]:
+    """Returns a clone of an existing function."""
+    f_new = FunctionType(
+        f.__code__,
+        f.__globals__,
+        name=f.__name__,
+        argdefs=f.__defaults__,
+        closure=f.__closure__,
+    )
+    f_new.__kwdefaults__ = f.__kwdefaults__
+    return update_wrapper(f_new, f)
+
+
 def lazy_xp_function(
-    func: Callable[..., Any],
+    func: Callable[..., Any] | tuple[type, str],
     *,
     allow_dask_compute: bool | int = False,
     jax_jit: bool = True,
@@ -69,8 +83,9 @@ def lazy_xp_function(
 
     Parameters
     ----------
-    func : callable
-        Function to be tested.
+    func : callable | tuple[type, str]
+        Function to be tested, or a tuple containing an (uninstantiated) class and a
+        method name to specify a class method to be tested.
     allow_dask_compute : bool | int, optional
         Whether `func` is allowed to internally materialize the Dask graph, or maximum
         number of times it is allowed to do so. This is typically triggered by
@@ -204,15 +219,49 @@ def lazy_xp_function(
             DeprecationWarning,
             stacklevel=2,
         )
-    tags = {
+    tags: dict[str, bool | int | type] = {
         "allow_dask_compute": allow_dask_compute,
         "jax_jit": jax_jit,
     }
 
+    if isinstance(func, tuple):
+        # Replace the method with a clone before adding tags
+        # to avoid adding unwanted tags to a parent method when
+        # the method was inherited from a parent class.
+        # Note: can't just accept an unbound method `cls.method_name` because in
+        # case of inheritance it would be impossible to attribute it to the child class.
+        # This also makes it so tagged methods will appear in their class's ``__dict__``
+        # and thus findable by ``iter_tagged_modules`` below.
+        cls, method_name = func
+        # The method might be a staticmethod or classmethod so we need to do a dance
+        # to ensure that this is preserved.
+        raw_attr = getattr_static(cls, method_name)
+        method = getattr(cls, method_name)
+        if isinstance(raw_attr, classmethod):
+            method = method.__func__
+        cloned_method = _clone_function(method)
+
+        method_to_set: Any
+        if isinstance(raw_attr, staticmethod):
+            method_to_set = staticmethod(cloned_method)
+        elif isinstance(raw_attr, classmethod):
+            method_to_set = classmethod(cloned_method)
+        else:
+            method_to_set = cloned_method
+
+        setattr(cls, method_name, method_to_set)
+        f = getattr(cls, method_name)
+        if isinstance(raw_attr, classmethod):
+            f = f.__func__
+        # Annotate that cls owns this method so we can check that later.
+        tags["owner"] = cls
+    else:
+        f = func
+
     try:
-        func._lazy_xp_function = tags  # type: ignore[attr-defined]  # pylint: disable=protected-access  # pyright: ignore[reportFunctionMemberAccess]
+        f._lazy_xp_function = tags  # pylint: disable=protected-access  # pyright: ignore[reportFunctionMemberAccess]
     except AttributeError:  # @cython.vectorize
-        _ufuncs_tags[func] = tags
+        _ufuncs_tags[f] = tags
 
 
 def patch_lazy_xp_functions(
@@ -224,10 +273,11 @@ def patch_lazy_xp_functions(
     """
     Test lazy execution of functions tagged with :func:`lazy_xp_function`.
 
-    If ``xp==jax.numpy``, search for all functions which have been tagged with
-    :func:`lazy_xp_function` in the globals of the module that defines the current test,
-    as well as in the ``lazy_xp_modules`` list in the globals of the same module,
-    and wrap them with :func:`jax.jit`. Unwrap them at the end of the test.
+    If ``xp==jax.numpy``, search for all functions and methods which have been tagged
+    with :func:`lazy_xp_function` in the globals of the module that defines the current
+    test, as well as in the ``lazy_xp_modules`` list in the globals of the same module,
+    and wrap them with :func:`jax.jit`.
+    Unwrap them at the end of the test.
 
     If ``xp==dask.array``, wrap the functions with a decorator that disables
     ``compute()`` and ``persist()`` and ensures that exceptions and warnings are raised
@@ -271,18 +321,34 @@ def patch_lazy_xp_functions(
     the example above.
     """
     mod = cast(ModuleType, request.module)
-    mods = [mod, *cast(list[ModuleType], getattr(mod, "lazy_xp_modules", []))]
-
-    to_revert: list[tuple[ModuleType, str, object]] = []
-
-    def temp_setattr(mod: ModuleType, name: str, func: object) -> None:
+    search_targets: list[ModuleType | type] = [
+        mod,
+        *cast(list[ModuleType], getattr(mod, "lazy_xp_modules", [])),
+    ]
+    # Also search for classes within the above modules which have had lazy_xp_function
+    # applied to methods through ``lazy_xp_function((cls, method_name))`` syntax.
+    # We might end up adding classes incidentally imported into modules, so using a
+    # set here to cut down on potential redundancy.
+    classes: set[type] = set()
+    for target in search_targets:
+        for obj in target.__dict__.values():
+            if isinstance(obj, type):
+                classes.add(obj)
+    search_targets.extend(classes)
+
+    to_revert: list[tuple[ModuleType | type, str, object]] = []
+
+    def temp_setattr(target: ModuleType | type, name: str, func: object) -> None:
         """
         Variant of monkeypatch.setattr, which allows monkey-patching only selected
         parameters of a test so that pytest-run-parallel can run on the remainder.
         """
-        assert hasattr(mod, name)
-        to_revert.append((mod, name, getattr(mod, name)))
-        setattr(mod, name, func)
+        assert hasattr(target, name)
+        # Need getattr_static because the attr could be a staticmethod or other
+        # descriptor and we don't want that to be stripped away.
+        original = getattr_static(target, name)
+        to_revert.append((target, name, original))
+        setattr(target, name, func)
 
     if monkeypatch is not None:
         warnings.warn(
@@ -298,10 +364,19 @@ def patch_lazy_xp_functions(
         temp_setattr = monkeypatch.setattr  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
 
     def iter_tagged() -> Iterator[
-        tuple[ModuleType, str, Callable[..., Any], dict[str, Any]]
+        tuple[ModuleType | type, str, Any, Callable[..., Any], dict[str, Any]]
     ]:
-        for mod in mods:
-            for name, func in mod.__dict__.items():
+        for target in search_targets:
+            for name, attr in target.__dict__.items():
+                # attr might be a staticmethod or classmethod. If so we need
+                # to peel it back and wrap the underlying function and later
+                # make sure not to accidentally replace it with a regular
+                # method.
+                func: Any = (
+                    attr.__func__
+                    if isinstance(attr, (staticmethod, classmethod))
+                    else attr
+                )
                 tags: dict[str, Any] | None = None
                 with contextlib.suppress(AttributeError):
                     tags = func._lazy_xp_function  # pylint: disable=protected-access
@@ -309,23 +384,49 @@ def patch_lazy_xp_functions(
                     with contextlib.suppress(KeyError, TypeError):
                         tags = _ufuncs_tags[func]
                 if tags is not None:
-                    yield mod, name, func, tags
-
+                    if isinstance(target, type) and tags.get("owner") is not target:
+                        # There's a common pattern to wrap functions in namespace
+                        # classes to bypass lazy_xp_function like this:
+                        #
+                        # class naked:
+                        #     myfunc = mymodule.myfunc
+                        #
+                        # To ensure this still works when checking for tags in
+                        # attributes of classes, ensure that target is the actual
+                        # owning class where func was defined.
+                        continue
+                    # put attr, and func in the outputs so we can later tell
+                    # if this was a staticmethod or classmethod.
+                    yield target, name, attr, func, tags
+
+    wrapped: Any
     if is_dask_namespace(xp):
-        for mod, name, func, tags in iter_tagged():
+        for target, name, attr, func, tags in iter_tagged():
             n = tags["allow_dask_compute"]
             if n is True:
                 n = 1_000_000
             elif n is False:
                 n = 0
             wrapped = _dask_wrap(func, n)
-            temp_setattr(mod, name, wrapped)
+            # If we're dealing with a staticmethod or classmethod, make
+            # sure things stay that way.
+            if isinstance(attr, staticmethod):
+                wrapped = staticmethod(wrapped)
+            elif isinstance(attr, classmethod):
+                wrapped = classmethod(wrapped)
+            temp_setattr(target, name, wrapped)
 
     elif is_jax_namespace(xp):
-        for mod, name, func, tags in iter_tagged():
+        for target, name, attr, func, tags in iter_tagged():
             if tags["jax_jit"]:
                 wrapped = jax_autojit(func)
-                temp_setattr(mod, name, wrapped)
+                # If we're dealing with a staticmethod or classmethod, make
+                # sure things stay that way.
+                if isinstance(attr, staticmethod):
+                    wrapped = staticmethod(wrapped)
+                elif isinstance(attr, classmethod):
+                    wrapped = classmethod(wrapped)
+                temp_setattr(target, name, wrapped)
 
     # We can't just decorate patch_lazy_xp_functions with
     # @contextlib.contextmanager because it would not work with the
@@ -335,8 +436,8 @@ def patch_lazy_xp_functions(
         try:
             yield
         finally:
-            for mod, name, orig_func in to_revert:
-                setattr(mod, name, orig_func)
+            for target, name, orig_func in to_revert:
+                setattr(target, name, orig_func)
 
     return revert_on_exit()
 

{array_api_extra-0.10.0 → array_api_extra-0.10.1}/tests/test_testing.py

@@ -1,9 +1,10 @@
 from collections.abc import Callable, Iterator
 from types import ModuleType
-from typing import cast
+from typing import Any, cast, final
 
 import numpy as np
 import pytest
+from typing_extensions import override
 
 from array_api_extra._lib._backends import Backend
 from array_api_extra._lib._testing import (
@@ -321,6 +322,128 @@ def test_lazy_xp_function_cython_ufuncs(xp: ModuleType, library: Backend):
         xp_assert_equal(cast(Array, erf(x)), xp.asarray([1.0, 1.0]))
 
 
+class A:
+    def __init__(self, x: Array):
+        xp = array_namespace(x)
+        self._xp: ModuleType = xp
+        self.x: Any = np.asarray(x)
+
+    def f(self, y: Array) -> Array:
+        return self._xp.asarray(np.matmul(self.x, np.asarray(y)))
+
+    def g(self, y: Array, z: Array) -> Array:
+        return self.f(y) + self.f(z)
+
+    def h(self, y: Array) -> bool:
+        return bool(self._xp.any(y))
+
+
+class B(A):
+    @override
+    def __init__(self, x: Array):  # pyright: ignore[reportMissingSuperCall]
+        xp = array_namespace(x)
+        self._xp: ModuleType = xp
+        self.x: Any = xp.asarray(x)
+
+    @override
+    def f(self, y: Array) -> Array:
+        return self._xp.matmul(self.x, y)
+
+    @staticmethod
+    def k(y: Array) -> "B":
+        return B(2.0 * y)
+
+    @staticmethod
+    def j(y: Array) -> "B":
+        xp = array_namespace(y)
+        y = xp.asarray(y)
+        if bool(xp.any(y)):
+            return B(y)
+        return B(y + 1.0)
+
+    @classmethod
+    def w(cls, y: Array) -> "B":
+        xp = array_namespace(y)
+        y = xp.asarray(y)
+        if bool(xp.any(y)):
+            return B(y)
+        return B(y + 1.0)
+
+
+@final
+class eager:
+    # this needs to be a staticmethod to appease the type checker
+    non_materializable5 = staticmethod(non_materializable5)
+
+
+lazy_xp_function((B, "g"))
+lazy_xp_function((B, "h"))
+lazy_xp_function((B, "k"))
+lazy_xp_function((B, "j"))
+lazy_xp_function((B, "w"))
+
+
+class TestLazyXpFunctionClasses:
+    def test_parent_method_not_tagged(self):
+        assert hasattr(B.g, "_lazy_xp_function")
+        assert not hasattr(A.g, "_lazy_xp_function")
+
+    @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="converts to NumPy")
+    @pytest.mark.skip_xp_backend(Backend.CUPY, reason="converts to NumPy")
+    @pytest.mark.skip_xp_backend(Backend.JAX_GPU, reason="converts to NumPy")
+    @pytest.mark.skip_xp_backend(Backend.TORCH_GPU, reason="converts to NumPy")
+    def test_lazy_xp_function_classes(self, xp: ModuleType, library: Backend):
+        x = xp.asarray([1.1, 2.2, 3.3])
+        y = xp.asarray([1.0, 2.0, 3.0])
+        foo = A(x)
+        bar = B(x)
+
+        if library.like(Backend.JAX):
+            with pytest.raises(
+                TypeError, match="Attempted boolean conversion of traced array"
+            ):
+                assert bar.h(y)
+
+        assert foo.h(y)
+
+    def test_static_methods_preserved(self, xp: ModuleType):
+        # Tests that static methods stay static methods when
+        # lazy_xp_function is applied.
+        x = xp.asarray([1.1, 2.2, 3.3])
+        foo = B(x)
+        bar = foo.k(x)
+        xp_assert_equal(bar.x, 2.0 * foo.x)
+
+    @pytest.mark.skip_xp_backend(Backend.DASK, reason="calls dask.compute()")
+    def test_static_methods_wrapped(self, xp: ModuleType, library: Backend):
+        x = xp.asarray([1.1, 2.2, 3.3])
+        foo = B(x)
+
+        if library.like(Backend.JAX):
+            with pytest.raises(
+                TypeError, match="Attempted boolean conversion of traced array"
+            ):
+                assert isinstance(foo.j(x), B)
+        else:
+            assert isinstance(foo.j(x), B)
+
+    @pytest.mark.skip_xp_backend(Backend.DASK, reason="calls dask.compute()")
+    def test_class_methods_wrapped(self, xp: ModuleType, library: Backend):
+        x = xp.asarray([1.1, 2.2, 3.3])
+        if library.like(Backend.JAX):
+            with pytest.raises(
+                TypeError, match="Attempted boolean conversion of traced array"
+            ):
+                assert isinstance(B.w(x), B)
+        else:
+            assert isinstance(B.w(x), B)
+
+    def test_circumvention(self, xp: ModuleType):
+        x = xp.asarray([1.0, 2.0])
+        y = eager.non_materializable5(x)
+        xp_assert_equal(y, x)
+
+
 def dask_raises(x: Array) -> Array:
     def _raises(x: Array) -> Array:
         # Test that map_blocks doesn't eagerly call the function;