ENH: Allow lazy_xp_function to work properly for class methods that are inherited from a parent class and for classmethods and staticmethods by steppi · Pull Request #582 · data-apis/array-api-extra

steppi · 2026-01-03T18:46:35Z

Closes #488

In #488 I asked about applying something like lazy_xp_function to classes. This was a half-baked idea motivated by wanting xp_capabilities in SciPy, which everyone here should be familiar with, to always be applied at the class level instead of the individual method level. In scipy/scipy#24267 I created a workflow that keeps xp_capabilities applied to classes, but still allows capabilities to be defined separately for individual methods if needed, and allows lazy_xp_function to be applied to individual methods separately. This means no nonsense about trying to automatically determine which methods of a class should have lazy_xp_function applied to them.

While working on this, I settled on specifying methods with tuples of the form Tuple[type, str] specifying an (uninstantiated) class and a method name. This is to allow distinguishing things like A.f from B.f when B is a subclass of A that inherits f from A, since capabilities may differ at different levels of the inheritance hierarchy. Through changes in SciPy, I was able to allow precise declarations of capabilities for class methods in the presence of inheritance, but a separate change is needed here to allow things like applying lazy_xp_function to B.f but not A.f when f is inherited from A.

The change here modifies lazy_xp_function to also take tuples of the form Tuple[type, str]. When this is done, for say (B, "f"), then B.f is replaced with a shallow clone of itself before adding the tags. This allowsB.f gets the tags without A.f to get them. If replacing an inherited method with a shallow clone is too obtrusive, I have a workaround that keeps more obtrusive modifications only within patch_lazy_xp_functions, but it makes things considerably more complicated, so I hope that won't be necessary.

https://github.com/scipy/scipy/blob/dfa1b87e4af7cf7ee5a8b8faf5c4360b63c86b36/scipy/_lib/tests/test_xp_capabilities.py has an (xfailed) test giving an example of what can go wrong with inheritance currently which can be made to pass with the change made in this PR.

lucascolley · 2026-01-03T19:29:20Z

a few CI failures otherwise looks pretty good

lucascolley · 2026-01-03T19:31:10Z

cc @crusaderky

steppi · 2026-01-03T21:19:10Z

Thanks @lucascolley. I've pushed a commit which fixes the failures locally.

lucascolley

thanks Albert LGTM

steppi · 2026-01-03T22:43:12Z

Coverage misses some lines in the diff because ufuncs allow setting attributes as of NumPy 2.2 (numpy/numpy#27735) so those lines only get hit for older NumPy versions.

lucascolley · 2026-01-03T22:51:33Z

Are they hit by pixi run -e tests-numpy1 tests-ci or not? If not, is there something else than just the NumPy version?

steppi · 2026-01-03T23:01:36Z

Are they hit by pixi run -e tests-numpy1 tests-ci or not? If not, is there something else than just the NumPy version?

Interesting, I tried it locally and it seems like they are still not hit. But they should be hit by this test if an older NumPy is installed::

try:
    # Test an arbitrary Cython ufunc (@cython.vectorize).
    # When SCIPY_ARRAY_API is not set, this is the same as
    # scipy.special.erf.
    from scipy.special._ufuncs import erf  # type: ignore[import-untyped]

    lazy_xp_function(erf)
except ImportError:
    erf = None


@pytest.mark.skip_xp_backend(Backend.TORCH_GPU, reason="device->host copy")
@pytest.mark.filterwarnings("ignore:__array_wrap__:DeprecationWarning")  # PyTorch
def test_lazy_xp_function_cython_ufuncs(xp: ModuleType, library: Backend):
    pytest.importorskip("scipy")
    assert erf is not None
    x = xp.asarray([6.0, 7.0])
    if library.like(Backend.ARRAY_API_STRICT, Backend.JAX):
        # array-api-strict arrays are auto-converted to NumPy
        # which results in an assertion error for mismatched namespaces
        # 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(cast(Array, erf(x)), xp.asarray([1.0, 1.0]))
    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(cast(Array, erf(x)), xp.asarray([1.0, 1.0]))

steppi · 2026-01-03T23:12:38Z

There's no mystery. The test that would hit those lines depends on SciPy and there's no SciPy in the tests-numpy1 env.

lucascolley · 2026-01-03T23:28:22Z

Wanna try pixi add -f tests scipy? If there are any difficulties then probably fine to just ignore it.

steppi · 2026-01-03T23:42:21Z

Wanna try pixi add -f tests scipy? If there are any difficulties then probably fine to just ignore it.

I can confirm that adding scipy causes those lines to get hit locally.

crusaderky

IMHO the API would feel more user-friendly if one could just write

lazy_xp_function(B.g)

is there a strong reason against it?

Also, is there a draft scipy PR that shows how this gets integrated in scipy's xp_capabilities?

steppi · 2026-01-07T13:57:58Z

Also, is there a draft scipy PR that shows how this gets integrated in scipy's xp_capabilities?

The draft PR is here scipy/scipy#24267.

there a strong reason against it?

In situations where B inherits g from A, B.g is a reference to the parent method and it is impossible to determine within the body of lazy_xp_function which class is intended from just B.g as far as I'm aware. Tuples allow unambiguous determination of which class is intended within the class hierarchy. This sidesteps issues that can arise when lazy capabilities differ at different levels of the hierarchy.

src/array_api_extra/testing.py

crusaderky · 2026-01-07T14:28:47Z

src/array_api_extra/testing.py

This is now obsolete. Could you either

clarify that you need to explicitly list classes as well as modules

change patch_lazy_xp_functions to descend into classes from the modules (preferrable)

Currently classes need to be explicitly listed as well as modules but yes, it would be nice if patch_lazy_xp_functions descended into modules so I'm +1 for that.

crusaderky · 2026-01-07T14:30:39Z

tests/test_testing.py

+def test_lazy_xp_function_class_inheritance():
+    assert hasattr(B.g, "_lazy_xp_function")
+    assert not hasattr(A.g, "_lazy_xp_function")


What this new test doesn't verify is that B.g actually runs with the JAX/Dask wrapper in a test with the xp fixture. In fact, it doesn't because patch_lazy_xp_functions doesn't descend into the class.

Yeah, because the class needs to be added to lazy_xp_modules and the test needs the xp fixture. Things will work in that case, and there are tests in SciPy that work after doing this. Maybe I should just update patch_lazy_xp_functions here too to sidestep that though.

crusaderky

.

steppi · 2026-01-07T18:45:47Z

Thanks @crusaderky, patch_lazy_xp_functions should correctly search classes within modules now.

src/array_api_extra/testing.py

crusaderky · 2026-01-08T15:56:25Z

tests/test_testing.py

+        foo = B(x)
+        observed = foo.g(y, z)
+        expected = xp.asarray(44.0)[()]
+        xp_assert_close(observed, expected)


This runs the function, but it doesn't test that it's been wrapped.
You need to write another function that will fail when wrapped, e.g.

def w(self): return bool(self._xp.any(self.x))`

See tests earlier in this same test module for examples.

steppi · 2026-01-08T17:21:41Z

Thanks @crusaderky; the test is actually doing what it's supposed to now.

lucascolley · 2026-01-08T17:30:48Z

tests/test_testing.py

+        foo = A(x)
+        bar = B(x)
+
+        if library.like(Backend.JAX):


not just JAX, sparse raises a RuntimeError, see CI

Yeah, but for different reasons orthogonal from what's being tested so I'm just going to add a skip for it. By the way, what's the pixi command to run tests with all backends?

array-api-extra on 🎋 lazy-class-methods is 📦 v0.10.0.dev0 via 🐍 took 4s ❯ pixi run tests ? The task 'tests' can be run in multiple environments. Please select an environment to run the task in: › ❯ tests tests-py313 dev tests-backends tests-backends-py311 dev-cuda tests-cuda tests-cuda-py311 tests-numpy1 tests-py311 tests-nogil

tests-backends for all (CPU) backends!

if in doubt the CI log will always show the exact Pixi task you need to reproduce

Thanks. Also, is there a convenient way to add a skip for all GPU backends? In SciPy skip_xp_backends has a cpu_only=True option, but here it looks like I'd need to add a skip for each GPU backend separately? I guess I could also just find a way to test the same behavior without needing the skips.

I don't think that exists yet, however it should be somewhat easy to add a method similar to

array-api-extra/src/array_api_extra/_lib/_backends.py

Lines 44 to 46 in b6518f3

def like(self, *others: Backend) -> bool: # numpydoc ignore=PR01,RT01

"""Check if this backend uses the same module as others."""

return any(self.modname == other.modname for other in others)

which checks if self.name ends with :gpu (or something more general).

Thanks. That's probably out of scope for this PR so I just added separate skips for all of the GPU backends.

steppi · 2026-01-09T01:18:32Z

I've just found that the change to make it so classes didn't have to be added to lazy_xp_modules broke some SciPy tests, so marking as draft until I have that figured out.

steppi · 2026-01-09T02:35:47Z

OK, so it wasn't actually related to the lazy_xp_modules stuff at all. The way I was cloning methods made it so staticmethods lost their staticmethodness. We need to do a dance to preserve it.

steppi · 2026-01-09T20:39:12Z

Marking as ready to review because CI is now fully green for scipy/scipy#24267.

Co-authored-by: Guido Imperiale <crusaderky@gmail.com>

lucascolley · 2026-01-09T21:48:19Z

the coverage for classmethod is probably not a big concern

steppi · 2026-01-09T21:50:48Z

the coverage for classmethod is probably not a big concern

I added tests for it locally and found that classmethods currently don't get wrapped still. There's another layer of indirection for classmethods compared to staticmethods. Working on fixing that now.

steppi · 2026-01-09T22:12:39Z

OK, that should do it.

lucascolley · 2026-01-09T22:44:10Z

the coverage for classmethod is probably not a big concern

I added tests for it locally and found that classmethods currently don't get wrapped still. There's another layer of indirection for classmethods compared to staticmethods. Working on fixing that now.

aha, wow I love codecov

rgommers · 2026-02-02T10:37:52Z

pyproject.toml

 hypothesis = ">=6.148.8"
 array-api-strict = ">=2.4.1,<2.5"
 numpy = ">=1.22.0"
+scipy = ">=1.15.2,<2"


This looks unnecessary/unwanted. Dropping this can then also drop the pixi.lock changes which will resolve the merge conflict.

It was deliberate for test coverage. I can handle conflicts

lucascolley added the enhancement New feature or request label Jan 3, 2026

lucascolley added the xpx.testing label Jan 3, 2026

lucascolley added this to the 0.10.0 milestone Jan 3, 2026

lucascolley approved these changes Jan 3, 2026

View reviewed changes

crusaderky reviewed Jan 7, 2026

View reviewed changes

src/array_api_extra/testing.py Show resolved Hide resolved

crusaderky approved these changes Jan 7, 2026

View reviewed changes

crusaderky reviewed Jan 7, 2026

View reviewed changes

crusaderky suggested changes Jan 7, 2026

View reviewed changes

lucascolley requested a review from crusaderky January 7, 2026 19:02

crusaderky suggested changes Jan 8, 2026

View reviewed changes

lucascolley requested a review from crusaderky January 8, 2026 17:23

lucascolley reviewed Jan 8, 2026

View reviewed changes

steppi marked this pull request as draft January 9, 2026 01:17

steppi marked this pull request as ready for review January 9, 2026 02:34

steppi marked this pull request as ready for review January 9, 2026 20:34

steppi and others added 18 commits January 9, 2026 21:06

ENH: Add better lazy_xp_function support for class methods

1224be7

TST: test that lazy_xp_function can handle inherited methods

5034b60

MAINT: Fix typing issues

c630610

ENH: make patch_lazy_xp_functions check classes within modules

4ba6e3e

TST: Add test to lazy_xp_function works for class methods

da76e14

Update src/array_api_extra/testing.py

96f7eda

Co-authored-by: Guido Imperiale <crusaderky@gmail.com>

Apply suggestions from code review

368e019

Co-authored-by: Guido Imperiale <crusaderky@gmail.com>

TST: make test laxy xp classes actually test function is wrapped

1e7997e

TST: Add appropriate skips for lazy xp function classes test

86cdec5

TST: Ensure that classmethods and staticmethods stay those things

35ccdd2

TST: Add test that lazy_xp_function preserves staticmethod

3401613

TST: Add test that static methods actually get wrapped

9107fe2

MAINT: fix monkeypatching of staticmethods and classmethods

2168009

TST: Make sure test_static_tests_wrapped doesn't fail

fd7d86f

MAINT: re-enable monkeypatch circumvention

6bdd40b

TST: add test that monkeypatching can be circumvented

9cf7895

MAINT: simplify check that target owns method

40dfb29

add scipy to test deps for coverage

0c12809

lucascolley force-pushed the lazy-class-methods branch from 5fec90b to 0c12809 Compare January 9, 2026 21:07

steppi added 2 commits January 9, 2026 17:11

TST: test that classmethods get wrapped

a739cfe

MAINT: make sure classmethods get wrapped

08b7e29

steppi changed the title ~~ENH: Allow lazy_xp_function to work properly for class methods that are inherited from a parent class~~ ENH: Allow lazy_xp_function to work properly for class methods that are inherited from a parent class and for classmethods and staticmethods Jan 9, 2026

Merge branch 'main' into lazy-class-methods

03da771

rgommers reviewed Feb 2, 2026

View reviewed changes

	def like(self, *others: Backend) -> bool: # numpydoc ignore=PR01,RT01
	"""Check if this backend uses the same module as others."""
	return any(self.modname == other.modname for other in others)

Conversation

steppi commented Jan 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lucascolley commented Jan 3, 2026

Uh oh!

lucascolley commented Jan 3, 2026

Uh oh!

steppi commented Jan 3, 2026

Uh oh!

lucascolley left a comment

Choose a reason for hiding this comment

Uh oh!

steppi commented Jan 3, 2026

Uh oh!

lucascolley commented Jan 3, 2026

Uh oh!

steppi commented Jan 3, 2026

Uh oh!

steppi commented Jan 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lucascolley commented Jan 3, 2026

Uh oh!

steppi commented Jan 3, 2026

Uh oh!

crusaderky left a comment

Choose a reason for hiding this comment

Uh oh!

steppi commented Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crusaderky left a comment

Choose a reason for hiding this comment

Uh oh!

steppi commented Jan 7, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

steppi commented Jan 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lucascolley Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

steppi commented Jan 9, 2026

Uh oh!

steppi commented Jan 9, 2026

Uh oh!

steppi commented Jan 9, 2026

Uh oh!

lucascolley commented Jan 9, 2026

Uh oh!

steppi commented Jan 9, 2026

Uh oh!

steppi commented Jan 3, 2026 •

edited

Loading

steppi commented Jan 3, 2026 •

edited

Loading

steppi commented Jan 7, 2026 •

edited

Loading

crusaderky Jan 7, 2026 •

edited

Loading

lucascolley Jan 8, 2026 •

edited

Loading