Useful functions not in the Array API Standard

[TomNicholas]

There are a few numpy functions which xarray calls on wrapped arrays but which are not (yet) in the Array API Standard. (See data-apis/array-api#187 (comment)) Cubed could choose to implement these to facilitate full integration.

FYI of this list xarray currently uses:

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

---

Of particular interest to me personally is np.pad. It's used within xarray's .pad method, which is used within xGCM's apply_as_grid_ufunc, which led to the pad function being an important part of the test case that exposed memory management problems with dask's distributed scheduler. I can't really close the loop by trying out cubed on that full original problem unless pad is available in cubed.

pad is also interesting because a parallel implementation isn't trivial - dask's pad implementation uses map_blocks in some cases, but more complicated tricks in other cases. For my purposes above I wouldn't need to implement more than one or two of the mode kwarg options though.

I would be interested in submitting a PR for adding pad if that's something you would welcome @tomwhite? (I mentioned this on the xarray tracker but it's really a cubed question pydata/xarray#7848 (comment))

cc @jbusecke

[tomwhite]

I would be interested in submitting a PR for adding pad

That would be very welcome! Just implementing the cases you need seems like a good way forward. (We might want to put it in a different namespace as it's not part of the array API, but we can discuss that later.)

BTW take is already implemented here.

[jbusecke]

That would be really cool @TomNicholas. Happy to test out prototypes whenever you think that is useful!

[tomwhite]

Xarray doesn't delegate to np.diff so this already works with Cubed:

>>> import xarray as xr
>>> import numpy as np
>>> import cubed.random
>>> da = xr.DataArray(cubed.random.random((3, 4), chunks=(2, 2)), dims=["x", "y"])
>>> d = da.diff("y")
>>> from numpy.testing import assert_array_equal
>>> assert_array_equal(d.values, np.diff(da.values, axis=1))

[tomwhite]

clip was added in #583

[tomwhite]

repeat was added in #610

[tomwhite]

diff was added to the 2024 standard in data-apis/array-api#791.

Also, pad has been added to the array-api-extra package in data-apis/array-api-extra#71 and data-apis/array-api-extra#72

[lucascolley]

to clarify, only the mode='constant' part of pad is added to array-api-extra at the moment

[tomwhite]

Thanks @lucascolley. In Cubed pad only supports symmetric, but I think we'll want constant too, to support shift in Xarray.