Support array_api_compat and array_api_extra API

[paugier]

From what I understood, the modern way of writing Python packages that use arrays is to use the Python array API standard instead of directly Numpy (see https://data-apis.org and also https://docs.scipy.org/doc/scipy-1.12.0/dev/api-dev/array_api.html).

Package maintainers should use https://data-apis.org/array-api-compat/ and https://data-apis.org/array-api-extra/ and write code like:

import array_api_extra as xpx
from array_api_compat import array_namespace

def compute(x):
    xp = array_namespace(x)
    y = xp.sum(x)
    ...
    return xpx.atleast_nd(y, ndim=2, xp=xp)

def your_function(x, y):
    xp = array_namespace(x, y)
    return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)

It would be useful if Pythran supported this new API.

Of course, Pythran would need to support only Numpy arrays, so xp would always be as np (see https://numpy.org/doc/stable/reference/array_api.html).

Moreover, I guess it would also be theoretically possible to write an implementation of the Python array API standard for Numpy arrays using Pythran 🙂. I guess it would be quite useless except maybe for testing (see https://data-apis.org/array-api/2024.12/verification_test_suite.html and https://github.com/data-apis/array-api-tests). I'm afraid it would be a bit long to compile.

[paugier]

I guess @rgommers could tell us if I understood correctly and what could be done in Pythran.

To be clearer, I was thinking that some array consuming libraries could use Pythran for some functions for Numpy arrays and fallback to Python for other array implementations. I would use that in packages like fluidsim, fluidimage and fluidfft.

It seems to me that it would be much easier if Pythran supported the array_namespace and array_api_extra API.

Then, there could also be a generic mechanism for the dispatch. A decorator is convenient for that so Transonic or something like Transonic could be used. (However, this last point is quite decoupled from what has to be done in Pythran.)

[rgommers]

Interesting idea. I'm not sure array-api-extra support is necessary (although more intelligence can always help of course). I think if the general pattern of xp = array_namespace(x) and then using xp instead of np is supported, that already makes it possible to write generic code.

That said, I'm not sure how much it will be used. The idea is then to support both a Python and a Pythran code path from the same code (a la scipy/scipy#22972), which could be helpful but only if other constructs that are used to support array API code are not getting in the way. As soon as there is some library-specific code (e.g., the if is_torch(xp) and if is_jax(xp) code paths in many scipy functions), it may get challenging to ignore/support for Pythran.

Maybe you have cleaner code with less backwards compatibility constraints in fluidsym et al. than we have in SciPy though - and then I can certainly see the appeal.

[paugier]

Ah I didn't realize that if is_torch(xp) and if is_jax(xp) were so important and common.

Note that with a tool like Transonic (which processes the code at build time to create backend specific functions), one could easily remove these code paths for backends that do not support them (for example Pythran).