ENH: jax_autojit

[crusaderky]

Supercharge xpx.testing.lazy_xp_function.

This is a fairly high level change; @rgommers @pearu @jakevdp I could use your feedback.

Matching SciPy PR: scipy/scipy#22909

Automatic static arguments

scipy has seen a proliferation of lazy_xp_function(func, static_argnames=(...)) in the initial section of many test modules. This information is only useful for JAX and is quite verbose.
With scipy/scipy#22686, this informational quirk that is both specific to JAX and to unit tests moves to the implementation modules.

With this PR, lazy_xp_function no longer accepts parameters static_argnames and static_argnums. Instead, all arguments that are not JAX arrays are automatically treated as static. Note that this behaviour is generally desirable in unit testing but a bad idea in production. Consider:

def f(x: Array, y: float, plus: bool) -> Array:
    return x + y if plus else x - y

j1 = jax.jit(f, static_argnames="plus")
j2 = jax_autojit(f)

jax_autojit is a new internal function of array-api-extra.
In the above example, j2 requires a lot less setup to be tested effectively, but on the flip side it means that it will be re-traced for every different value of y, which likely makes it not fit for purpose in production.

To clarify: jax_autojit is applied and removed on the fly within tests with the xp fixture and it is never used outside of unit testing.

Wrapped inputs and output

There are a few cases of scipy functions returning bespoke containers with arrays inside instead of simple tuples, namedtuples, lists, or dicts of arrays.

Two such examples are

• scipy.stats.ttest_ind
• scipy.spatial.Rotation and all functions and methods that return it (discussion: ENH: spatial.transform: add array API standard support scipy/scipy#22777 (comment))

In main, these can't be tested with lazy_xp_function, and as of scipy/scipy#22686 the issue also impacts documentation.

This PR lifts this restriction and allows completely arbitrary objects as parameters and as return values of the functions. If these objects internally contain JAX arrays, lazy_xp_function will now automatically extract them, pass them through the JIT, and reassemble everything for the test function to observe the result.

The rationale is that, in real life, users are unlikely to wrap the scipy functions with jax.jit directly; instead they are more likely to consume their outputs in their own functions and then wrap those with jax.jit:

@jax.jit
def my_user_function(x):
    y = scipy.stats.ttest_ind(x)
    return my_user_consume(y)  # returns array or tuple of arrays

Static non-hashable arguments

Non-hashable objects can now be static.
Consider:

def f(x: Array, verbs: Sequence[str]) -> Array:
    ...
    
lazy_xp_function(f, static_argnames="verbs")

def test_f(xp):
    x = ...
    y = f(x, ("foo", "bar", "baz"))  # OK
    y = f(x, ["foo", "bar", "baz"])  # Fails

This PR fixes it; all objects in the input and output of the function need only be hashable or pickleable.

Dask materialization raises inside a container

This also fixes a Dask-specific bug where the graph materialization raises:

@pytest.mark.skip_xp_backends("jax.numpy", reason="raise inside jax.pure_callback")
def test1(xp):
    with pytest.raises(Exception):
        f(x)

The above works when f returns plain Dask array objects or tuples or lists thereof, even if the graph would otherwise be discarded and never computed, thanks to the lazy_xp_function machinery, but used to fail when the return value is an opaque container with Dask arrays inside. This PR fixes it.

[crusaderky]

Aside note on design: Dask avoids this exact problem by not requiring any decorator and instead duck-type checking for uniquely named dunder methods called __dask_<...>__

[rgommers]

Interesting design tradeoffs there - to me the "eager by default, opt-in for graph mode" is nicer and has won in array land though (dataframes are a different story). I guess this code pattern is one of the prices to pay.

[crusaderky]

Note: everything in this module is a private helper.

[pearu]

I have a minor nit and a note, otherwise it looks good to me as it reduces the maintenance burden of adding/updating static_argnames/argnums arguments.

Btw, I enjoyed reading pickle_flatten and pickle_flatten implementations as it felt like less is more.

Thanks, @crusaderky!

[jakevdp]

For what it's worth, one of the main reasons we haven't implemented something like this in JAX is because it tends to negatively impact dispatch times. The output of jax.jit is a C++ level callable, that directly dispatches to the compiled kernel after the initial call. This is important because if you put a layer of Python logic in the dispatch path to inspect all the arguments and determine whether to treat them as static or not (and possibly re-compile based on this information) then it severely impacts dispatch times.

With that background, I'd like a bit of clarification here: is this intended for use only in testing paths, or do you imagine this will be used within the dispatch path for libraries like scipy that implement the Array API?

[lucascolley]

I think this section of the top post answers your question, Jake:

To clarify: jax_autojit is applied and removed on the fly within tests with the xp fixture and it is never used outside of unit testing.

[crusaderky]

scipy/scipy#22909 does not cause any issues to crop up.
This PR is ready for final review.

[lucascolley]

@crusaderky will this close gh-270?

[crusaderky]

@crusaderky will this close gh-270?

Yes, it does!

[rgommers]

@lucascolley are you happy with me merging this once I'm done reviewing, or do you want to review it as well? I'd like to get it into scipy 1.16.x if possible. This PR itself isn't quite user-facing, but the docs improvements that build on this PR are.

[lucascolley]

Yes, I took a look through from a non-expert perspective and nothing seems untoward!

[rgommers]

This is very nice, it makes testing of JAX more generic and much easier to understand for maintainers and new contributors to SciPy et al. who don't have much (or any) JAX experience. The test simplifications in this PR show big the difference is.

👍🏼 for the PR description as well, that was very useful to me - and I expect others are going to refer to it in the future.

Btw, I enjoyed reading pickle_flatten and pickle_flatten implementations as it felt like less is more.

💯 agreed

I have some small non-blocking doc suggestions only. Those can wait or be left alone; I'll go ahead and merge this PR now so we call pull it in to unblock the SciPy release.

The machinery added here is quite complex, so it's very well possible I missed something when reviewing. I tested fairly extensively with SciPy and all looks good though, so nothing major should be off here.

Thanks @crusaderky & reviewers!

[rgommers]

Interesting design tradeoffs there - to me the "eager by default, opt-in for graph mode" is nicer and has won in array land though (dataframes are a different story). I guess this code pattern is one of the prices to pay.

[rgommers]

These seem like a good set of choices for testing purposes to me. For the unsuspecting reader who may be looking at this code without much context, it'd be useful to add something like this:

Note: these are useful choices *for testing purposes only*, which is how this function is
intended to be used.

[rgommers]

minor: this description starting with "Set to True" hints at the default being False, which is not the case. I'd switch the True/False cases around and start the True case with "When set to True (default)"

Note that this isn't a new issue, it just gets magnified because the "Default: True is a lot further down".