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