JIT?

[34j]

I am not sure how to define a JIT decorator which support multiple libraries. Maybe like this?

from collections.abc import Mapping, Sequence
from functools import cache, wraps
from types import ModuleType
from typing import Any, Callable, ParamSpec, TypeVar

from array_api_compat import array_namespace
from frozendict import frozendict

P = ParamSpec("P")
T = TypeVar("T")


def jit_impl(
    f: Callable[P, T], module: ModuleType, /, *args: Any, **kwargs: Any
) -> Callable[P, T]:
    if "numpy" in module.__name__:
        import numba

        return numba.jit(*args, **kwargs)(f)
    elif "torch" in module.__name__:
        import torch

        return torch.jit.script(f, *args, **kwargs)
    else:
        return getattr(module, "jit", lambda x: x)(*args, **kwargs)(f)


def jit(
    *,
    args: Mapping[ModuleType, Sequence[Any]] | None = None,
    kwargs: Mapping[ModuleType, Mapping[str, Any]] | None = None,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    def decorator(f: Callable[P, T]) -> Callable[P, T]:
        args_ = frozendict(args or {})
        kwargs_ = frozendict(kwargs or {})
        jit_cached = cache(jit_impl)

        @wraps(f)
        def inner(*args_inner: P.args, **kwargs_inner: P.kwargs) -> T:
            try:
                xp = array_namespace(*args_inner)
            except TypeError as e:
                if e.args[0] == "Unrecognized array input":
                    return f(*args_inner, **kwargs_inner)
                raise
            f_jit = jit_cached(f, xp, *args_.get(xp, ()), **kwargs_.get(xp, {}))
            return f_jit(*args_inner, **kwargs_inner)

        return inner
    return decorator

[lucascolley]

cc @crusaderky

[34j]

I cannot even compile a function because of array_namespace()...

E           torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
E             File ".venv/lib/python3.13/site-packages/array_api_compat/common/_helpers.py", line 489
E           def array_namespace(
E               *xs: Array | complex | None,
E               ~~~ <--- HERE
E               api_version: str | None = None,
E               use_compat: bool | None = None,

If I manually assign array_api_compat.torch inside the wrapped function, then it will be the genuine torch for unknown reason.

EDIT: torch.compile worked!

[34j]

E           numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
E           Untyped global name 'array_namespace': Cannot determine Numba type of <class 'function'>
E           
E           File "tests/test_main.py", line 7:
E           def spherical_coordinates(r: Any, theta: Any, phi: Any) -> Any:
E               xp = array_namespace(r, theta, phi)
E               ^
E           
E           During: Pass nopython_type_inference

.venv/lib/python3.13/site-packages/numba/core/dispatcher.py:365: TypingError

[34j]

Then I can't think of any other way to use JIT-like something with array-api-compat & torch but to use closed source ivy.trace_graph().

[crusaderky]

I am not sure how to define a JIT decorator which support multiple libraries.

Don't.

There is enough nuance in the keywords needed by each jit function that it's territory I really would not want to venture in.

A package that supports multiple Array API compatible libraries

• should test its code by wrapping it, within the scope of unit tests only, with the jit decorator of different libraries (e.g. JAX, torch)
• should not offer an "auto-jit" feature that somehow, maybe, kind of, works with everything out of the box.

For example, jax.jit is a net loss in terms of performance if you're only going to call each function once, change the shape of your array, and do it again. JAX does offer a way to move away from shape and to pin the tracing to ndim instead, where you basically promise explicitly that the number of elements along this or that axis won't change the outcome of your pure-python code. That is frequently non-trivial to guarantee, as it is easy for example to have special code paths for size 0 or size 1.

Another point: jax.jit performance for small-ish arrays relies on the fact that, once the jit is warm, there is no python involved whatsoever. A generic wrapper that identifies the input and auto-applies different jit algorithms would destroy this.

I am not overly familiar with torch.compile, but I expect the same amount of different nuance for it.

You can read an informative discussion here: #284

[crusaderky]

E numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)

That's not how numba.jit works.

jax.jit runs completely arbitrary pure python code and builds a graph from it by hot-swapping the real arrays in the inputs of your functions with tracers. Conceptually, it's the same as Dask.

numba.jit reads your code and tries to transpile it into LLVM. As a result, it supports a very, very limited subset of what Python can do. It only works with code that you specifically wrote for it, and is designed as an easy translation for C algorithms.

[34j]

But if numpy is array-api compatible, numba should work by like removing xp = array_api_namespace and replacing xp with np

[crusaderky]

No, this is not how numba works. A very tiny subset of the current Python code that invokes numpy API will let you do that.

[34j]

#360 (comment) will allow to use array-api compatible part of Numpy to be JIT-ed while making the code still compatible with other libraries

[34j]

Numba

@overload allowed numba to "ignore" array_api_compat.array_namespace():

Attempt 1

import numba
import numpy as np
from array_api_compat import array_namespace


@numba.njit
def sinplus1(x):
    xp = array_namespace(x)
    return xp.sin(x) + 1


print(sinplus1(np.arange(10)))

Traceback (most recent call last):
  File "test.py", line 12, in <module>
    sinplus1(np.arange(1000))
    ~~~~~~~~^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.13/site-packages/numba/core/dispatcher.py", line 424, in _compile_for_args
    error_rewrite(e, 'typing')
    ~~~~~~~~~~~~~^^^^^^^^^^^^^
  File ".venv/lib/python3.13/site-packages/numba/core/dispatcher.py", line 365, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'array_namespace': Cannot determine Numba type of <class 'function'>

File "test.py", line 8:
def sinplus1(x):
    xp = array_namespace(x)
    ^

During: Pass nopython_type_inference

Attempt 2

import numba
import numpy as np
from array_api_compat import array_namespace
from numba.extending import overload


@overload(array_namespace)
def array_namespace_overload(*args):
    def inner(*args):
        return np

    return inner


@numba.njit
def sinplus1(x):
    xp = array_namespace(x)
    return xp.sin(x) + 1


print(sinplus1(np.arange(10)))

[1.         1.84147098 1.90929743 1.14112001 0.2431975  0.04107573
 0.7205845  1.6569866  1.98935825 1.41211849]

Attempt 3

(Add **kwargs)

InternalError: The use of VAR_KEYWORD (e.g. **kwargs) is unsupported. (offending argument name is '**kwargs')

Attempt 4

(Add api_version=None, use_compat=None explicitly)

IndexError: tuple index out of range

numba/numba#5291
numba/numba#8108
numba/numba#8904

[crusaderky]

Before you start figuring out how to convince numba to ignore arbitrary stuff, I implore you, stop and first find a use case that is faster in numba than it is in plain numpy, while remaining performant with plain numpy (so, fully vectorized and using Array API patterns).
Because you will not find one.