RFC: `xpx.compile` as an abstraction for `torch.compile` and `jax.jit` decorators?

[ogrisel]

Has anybody investigated the possibility to allow for an array agnostic way to leverage the torch.compile and jax.jit decorators in array-api-extra?

This might be useful for array API consuming libraries such as SciPy or scikit-learn. For array API namespaces without JIT compiler support, xpx.compile would just result in a noop decorator. For torch and JAX it might, dispatching to an actual JIT compiler could unlock significant speed-ups and memory usage improvements.

However, the parameters of those decorators have many kwargs with seemingly very little overlap:

• https://docs.jax.dev/en/latest/_autosummary/jax.jit.html
• https://docs.pytorch.org/docs/stable/generated/torch.compile.html

Maybe xpx.compile could be made to accept arbitrary kwargs scoped by the underlying namespace name without attempting to map common compiler semantics together.

@xpx.compile(
   torch=dict(options={"triton.cudagraphs": True}, fullgraph=True),
   jax=dict(static_argnames=['n']),
)
def some_array_function(array, n):
   ...

I have little experience to tell whether calling those decorators with their default argument is useful or not in practice.

[lucascolley]

cc @crusaderky

[crusaderky]

As already discussed with JAX maintainers in #284, this would introduce a major overhead to JAX dispatch times, as you would need, on every invocation, to scan all args and kwargs, also descending in lists and dicts, to see if there are any JAX or torch arrays to be found.

[ogrisel]

Thanks, I missed the jax_autojit feature when searching for previous compiler related discussions in this repo before opening this issue.

Indeed, such as namespace agnostic decorator will require to rescan the namespace of the inputs at each call to know which namespace-specific compiler decorator to dispatch to. But I don't see a way around if we want to make it possible to keep namespace agnosticism in an array API consuming library.

[ev-br]

cross-ref a JIT discussion in SciPy: scipy/scipy#23447

[purepani]

This might be useful for array API consuming libraries such as SciPy or scikit-learn.

I think the pattern is generally to have the compilation be done at the last step by the end user, particularly because of how different it is between array libraries. I.e. as long as you can confirm that your functions are compatible with jit, the end user can jit compile it if they want to.

[ogrisel]

I.e. as long as you can confirm that your functions are compatible with jit, the end user can jit compile it if they want to.

Unfortunately, for a library like scikit-learn, most of the public API exposed to the user is unlikely to be jittable. I suspect that the jittable parts of scikit-learn are mostly private functions and methods.

EDIT: the way jax-sklearn is designed seems to confirm that: scikit-learn/scikit-learn#29647 (comment).

[ev-br]

cross-ref a long discussion of a related issue in scipy: scipy/scipy#23447
(The particular scipy use case is very similar to scikit-learn estimators)