Compatibility with the Array API standard

[TomNicholas]

What is your issue?

Meta-issue to track all the smaller issues around making xarray and the array API standard compatible with each other.

We've already had

• Support NumPy array API (experimental) #6804
• More Array API changes #7067
• Array API fixes for astype #7847

and there will likely be many others.

---

I suspect this might require changes to the standard as well as to xarray - in particular see this list of common numpy functions which are not currently in the array API standard. Of these xarray currently uses (FYI @ralfgommers ):

• np.clip
• np.diff
• np.pad
• np.repeat
• np.take
• np.tile

[TomNicholas]

np.pad is an interesting example in the context of chunked arrays (xref #6807) - dask implements it (parallelized using various approaches, including map_blocks internally), but cubed currently doesn't implement it because it's not part of the array API standard. (cc @tomwhite)

[shoyer]

Probably the most important missing feature for Xarray in the Array API standard is support for indexing with integer arrays: data-apis/array-api#669

[kgryte]

Quick update from the Array API standard, both take and tile have been merged in the standard. take is available as of the 2022.12 revision and tile will be included in the 2023 revision.

[TomNicholas]

Thanks @kgryte! We also just bumped our minimum-supported numpy to 1.22 (#8346), which means we could potentially start using numpy.array_api internally for various operations...

[allen-adastra]

Hi! I'm wondering if this is still a concerted effort? Def would love to see tight Xarray <-> Jax integration.

[keewis]

I don't have any knowledge about jax specifically, but packages like cubed have been wrapped successfully by xarray using the existing Array API support. So I'd say that as long as jax fully supports the Array API, you should be able to do this already. If not, could you please open a dedicated issue?

There might be a few operations that are needlessly verbose, though, so someone would have to create a integration package (jax-xarray or xarray-jax, not sure).

[TomNicholas]

Hey @allen-adastra, great to see you here.

Wrapping JAX with Xarray is a great idea in need of a champion / beta-tester. As Justus says it should just work, but most of the Xarray devs concentrate on making the high-level wrappers compliant, rather than exposing all the edge cases the specific libraries.

If you or one of your colleagues wants to just have a go and report back what happens that would be a great start!

[allen-adastra]

Gotcha, so it sounds like Xarray should have compliance with the Array API?

Probably too much for me to take on, but going to talk around to see if we can drum up sufficient support somehow...