Namespace-aware xarray.ufuncs

[slevang]

• Closes Compatibility with the Array API standard  #7848 (partially)
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

Re-implement the old xarray.ufuncs module to allow generic ufunc handling for array types that don't implement __array_ufunc__:

import jax.numpy as jnp
import numpy as np
import xarray as xr
import xarray.ufuncs as xu

x = xr.DataArray(jnp.asarray([1, 2, 3]))
print(type(xu.sin(x).data))
print(type(np.sin(x).data))

# <class 'jaxlib.xla_extension.ArrayImpl'>
# <class 'numpy.ndarray'>

[slevang]

Should we ever prioritize dispatching with np.func via __array_ufunc__ (if it exists) over the library's __array_namespace__().func?

[slevang]

Is there ever a reason to use dask's ufuncs with dask="allowed" instead of the appropriate _meta array's namespace and dask="parallelized"? With jax for example, which doesn't have __array_ufunc__, this ends up converting to numpy. So it would have to be special cased.

[TomNicholas]

In user code using xr.apply_ufunc there is - dask='allowed' can be used to rechunk along a core dimension e.g. by applying a dask reduction ufunc along that dimension. Not sure if that's relevant here though.

[slevang]

These are all elementwise so no core dimensions

[slevang]

Can hard code these if preferred?

[slevang]

These are the ones that didn't immediately work. There are also other ufunc like things that aren't technically np.ufunc subclasses that we could add. I saw angle and iscomplex were special cased before.