Compatibility with the Array API standard

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

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~~

@tomwhite

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)

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

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.

@kgryte

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...

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

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).

@allen-adastra

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!

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...

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

at least as far as dispatching to wrapped arrays goes, yes. The full Array API is not a good fit for xarray's objects and adopting just a subset is not allowed, so we don't implement it for those.

@keewis

To clarify, xarray aims to be able to wrap any array-API compliant object. This should already work, and is already being used to wrap some array types successfully (e.g. cubed), but it's likely that trying this with JAX will turn up some corner cases.

@keewis I think is referring to the separate idea of xarray objects (e.g. xr.Variable) following the array API, (as wrappable objects rather than wrappers), which we decided against.

So all I'm suggesting that you guys do @allen-adastra is try wrapping a JAX array with xarray and see what happens, and report back any failures. I'm not suggesting you take on the task of making xarray work for any array-API compliant object (as we're already trying to do that for the general case).

Now I think of it: this is not really xarray-specific, but I would have expected there to be a canonical 'generic' array API namespace that dispatches based on array_namespace of the arguments, e.g.
import array_api.generic_ufuncs as xnp

xnp.exp(jax_array) -> jax_array
xnp.exp(numpy_array) -> numpy_array
Does anything like this exist, and if not, should it exist? and could xarray piggy-pack off it without implementing the array API for xarray datatypes directly?

The problem is that Xarray objects shouldn't implement __array_namespace__, because they are not fully array API compliant. We talked about this pretty extensively in data-apis/array-api#698

If anyone wants to take a stab at adding this xarray.ufuncs module it would just have to be defined at the top-level and then examine the wrapped array type .data to decide which namespace to use on it for the particular operation.

I would have a go at this, although it looks like it would be easier for someone with more context about the old ufuncs.py. The old version contains some multi-argument dispatch logic which I'm not sure if you'd want to preserve or not, for example. And it goes via _unary_op and _binary_op, so not clear to me if the dispatch based on underlying array type should live there, or in ufuncs.py. If it lives in _unary_op/_binary_op, could that be disentangled from the dispatch logic in ufuncs.py, etc.

Alternatively would ufuncs.py now be better implemented just as thin wrappers around apply_ufunc?

Alternatively would ufuncs.py now be better implemented just as thin wrappers around apply_ufunc?

Yes, I would lean towards this. If I recall correctly, xarray.ufuncs pre-dated apply_ufunc.

Add `__array_namespace__` to `ArrayBox` (#647)

I ran into the ufuncs issue trying to do basic ops on objects containing jax arrays. Opened #9776 for discussion.

might be good to reopen this for further discussion

Opened #9798 as another follow up, curious to get feedback.

I was wondering what the status is of this issue. I am using a package (pyet meteo_utils.py::326) that calls numpy.clip. With the combination of numpy 2.1.3 and xarray 2025.1.2, I get an error for TypeError: DataWithCoords.clip() got an unexpected keyword argument 'out'.

The best solution would be that xarray would be compatible with numpy.clip. An alternative would be that I make a PR for that 3rd party package to solve it.

Is there any chance that the last functions are going to be compatible with numpy soon?

I was wondering what the status is of this issue. I am using a package (pyet meteo_utils.py::326) that calls numpy.clip. With the combination of numpy 2.1.3 and xarray 2025.1.2, I get an error for TypeError: DataWithCoords.clip() got an unexpected keyword argument 'out'.

This looks like a different NumPy/Xarray compatibility quirk, likely unrelated to the Array API standard. Could you please open another issue to discuss?

I was wondering what the status is of this issue. I am using a package (pyet meteo_utils.py::326) that calls numpy.clip. With the combination of numpy 2.1.3 and xarray 2025.1.2, I get an error for TypeError: DataWithCoords.clip() got an unexpected keyword argument 'out'.

This looks like a different NumPy/Xarray compatibility quirk, likely unrelated to the Array API standard. Could you please open another issue to discuss?

This was all my bad. The issue occurs due to a DeprecationWarning and only occurs in pytest. I was not able to reproduce it without pytest.

[pytest]
addopts = --ff 
testpaths = .

filterwarnings =
  error

This ensures that DeprecationWarnings are treated as errors, and therefore I get the error. I guess this still needs to be fixed at a later stage. But I'm leaving it for now.

Minimal example

import xarray as xr
import numpy as np


def test_clip():
    da = xr.DataArray(
        np.random.rand(4, 3),
        dims=("x", "y"),
        coords={"x": [0, 1, 2, 3], "y": [0, 1, 2]},
    )
    clipped = np.clip(da, 0.3, 0.6)

I can reproduce a DeprecationWarning with the following, but it is different to the warning quoted above:

❯ pixi exec -s xarray -s ipython -s numpy=2.1.3 ipython
<snip>
In [1]: import warnings

In [2]: warnings.filterwarnings('always')

In [3]: import xarray as xr
   ...: import numpy as np
   ...: 
   ...: da = xr.DataArray(
   ...:     np.random.rand(4, 3),
   ...:     dims=("x", "y"),
   ...:     coords={"x": [0, 1, 2, 3], "y": [0, 1, 2]},
   ...: )
   ...: np.clip(da, 0.3, 0.6)
<snip>/numpy/_core/fromnumeric.py:48: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)
  return conv.wrap(result, to_scalar=False)
Out[3]: 
<xarray.DataArray (x: 4, y: 3)> Size: 96B
<snip>

In [4]: np.__version__
Out[4]: '2.1.3'

In [5]: xr.__version__
Out[5]: '2025.1.2'

I see the same warning with latest NumPy.

Can we please fork the deprecation warning about np.clip to another issue? It has nothing to do with the array API standard, besides the fact that they both involve a function named clip.

Navigation Menu

Repository navigation

Compatibility with the Array API standard #7848

What is your issue?

13 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Compatibility with the Array API standard #7848

Description

What is your issue?

Activity

TomNicholas commented on May 18, 2023

shoyer commented on Aug 14, 2023

kgryte commented on Oct 19, 2023

TomNicholas commented on Oct 19, 2023

allen-adastra commented on May 19, 2024

keewis commented on May 19, 2024

TomNicholas commented on May 19, 2024

allen-adastra commented on May 19, 2024

keewis commented on May 19, 2024

TomNicholas commented on May 20, 2024

13 remaining items

shoyer commented on Sep 10, 2024

mjwillson commented on Sep 16, 2024

mjwillson commented on Sep 17, 2024

shoyer commented on Sep 17, 2024

slevang commented on Nov 13, 2024

lucascolley commented on Nov 18, 2024

slevang commented on Nov 18, 2024

deltamarnix commented on Mar 6, 2025

shoyer commented on Mar 6, 2025

deltamarnix commented on Mar 6, 2025

lucascolley commented on Mar 6, 2025

shoyer commented on Mar 6, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions