Guidance for creating Xarray consumer libraries (supporting flexible array backends, and in general)

[VeckoTheGecko]

Heyall . I have been wanting to make Parcels as an Xarray consuming library be flexible to the array backend (i.e., if we're doing operations it would be great to be able to avoid - for example - calling np.sum on dask/cupy arrays triggering eager computation or communication between the GPU and CPU.

This is a bit of thinking ahead for us as we want to avoid unnecessary Dask computations, support GPU in future (via Cupy), and just generally allow our support for array backends to evolve alongside Xarray.

What would be the best way to do this in Xarray world? I attended the talk from @lucascolley at EuroScipy 2025 talking about the Array API ecosystem - of particular interest being the array-api-compat library. From my understanding though, this isn't the full picture - partly at least because Dask is not array API compliant dask/dask#10994 (nor is cubed for that matter).

I see that there is a xarray.compat.array_api_compat (a custom file providing an array API compat layer - not to be confused with the array-api-compat package) providing namespacing for Array API compliant packages, but there is still xarray/core/duck_array_ops.py to deal with different array types (e.g., for Dask). Hence the following - which I would have expected to return da where Dask - doesn't work

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "cubed",
#     "xarray[complete]",
# ]
# ///
from xarray.compat.array_api_compat import get_array_namespace
import dask.array as da
import numpy as np
import xarray as xr

# breakpoint()
for data_array, name in [
    (xr.DataArray(da.array([1, 2, 3])), "Dask"),
    (xr.DataArray(np.array([1, 2, 3])), "numpy"),
]:
    xp = get_array_namespace(data_array)
    print(f"Namespace for DataArray backed by {name} array: ", str(xp))

    xp = get_array_namespace(data_array.values)
    print(f"Namespace for DataArray.values backed by {name} array: ", str(xp))

"""Output:
Namespace for DataArray backed by Dask array:  <module 'numpy' from '/Users/Hodgs004/.cache/uv/environments-v2/test-c07785d475a4d14f/lib/python3.12/site-packages/numpy/__init__.py'>
Namespace for DataArray.values backed by Dask array:  <module 'numpy' from '/Users/Hodgs004/.cache/uv/environments-v2/test-c07785d475a4d14f/lib/python3.12/site-packages/numpy/__init__.py'>
Namespace for DataArray backed by numpy array:  <module 'numpy' from '/Users/Hodgs004/.cache/uv/environments-v2/test-c07785d475a4d14f/lib/python3.12/site-packages/numpy/__init__.py'>
Namespace for DataArray.values backed by numpy array:  <module 'numpy' from '/Users/Hodgs004/.cache/uv/environments-v2/test-c07785d475a4d14f/lib/python3.12/site-packages/numpy/__init__.py'>
"""

So my question is the following:

What is the current approach for making and testing an Xarray consuming library to be flexible to array backend ? Is there a guide on this? Are there such consuming libraries code I can poke around in more to see how they handle things? (tried xdggs, but couldn't find this code) Regarding testing, it would be great to make sure that arrays don't change type unexpectedly

More generally, is there documentation with what considerations to have when creating xarray consuming libraries? (I guess somewhat related is the docs on accessors, are there other things as well?)

[tomwhite]

We had a similar challenge in sgkit - how to make sgkit's xarray usage support different parallel backends (Dask and Cubed). The approach we took is summarised here: sgkit-dev/sgkit#908 (comment), but basically

1. move from dask.array.map_blocks to xarray.apply_ufunc for applying functions in parallel,
2. have a way to run the test suite using either Dask or Cubed so we don't have to do all the changes at once.

I attended the talk from @lucascolley at EuroScipy 2025 talking about the Array API ecosystem - of particular interest being the array-api-compat library. From my understanding though, this isn't the full picture - partly at least because Dask is not array API compliant dask/dask#10994 (nor is cubed for that matter).

Using Dask via array-api-compat should be array API compliant though - did this not work for you? Also Cubed is array API compliant too, except for some gaps documented at https://github.com/cubed-dev/cubed/blob/main/api_status.md.

[lucascolley]

Have you considered contributing the corresponding https://data-apis.org/array-api-compat/helper-functions.html for Cubed to array-api-compat btw @tomwhite ?

[tomwhite]

No I hadn't - thanks for the suggestion!