Improved duck array wrapping #9798

slevang · 2024-11-18T20:55:46Z

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

Companion to #9776.

My attempts to use xarray-wrapped jax arrays turned up a bunch of limitations of our duck array wrapping. Jax is probably worst case of the major duck array types out there, because it doesn't implement either __array_function__ or __array_ufunc__ to intercept numpy calls. Cupy, sparse, and probably others do, so calling np.func(cp.ndarray) generally works fine, but this usually converts your data to numpy with jax.

A lot of this was just grepping around for hard-coded np. cases that we can easily dispatch to duck_array_ops versions. I image some of these changes could be controversial, because a number of them (notably nanmean/std/var, pad, quantile, einsum, cross) aren't technically part of the standard. See #8834 (comment) for discussion about the nan-skipping aggregations.

It feels like a much better user experience though to try our best to dispatch to the correct backend, and error if the function isn't implemented, rather than blindly calling numpy. And practically, all major array backends that are feasible to wrap today (cupy, sparse, jax, cubed, arkouda, ...?) implement all of these functions.

To test, I just ran down the API list and ran most functions to see if we maintain proper wrapping. Prior to the changes here, I had 28 jax failures and 9 cupy failures, while all (non-xfailed) ones now pass.

Basically everything works except the interp/missing methods which have a lot of specialized code. Also a few odds and ends like polyfit and rank.

slevang · 2024-11-18T20:58:25Z

xarray/core/array_api_compat.py

+        if hasattr(x, "__array_namespace__"):
+            return x.__array_namespace__()
+        elif isinstance(x, array_type("cupy")):
+            # special case cupy for now


cupy seems to have full compliance with the standard, but doesn't yet actually have __array_namespace__ on the core API. Others may be the same?

slevang · 2024-11-18T20:59:17Z

xarray/core/computation.py

@@ -2174,9 +2174,13 @@ def _calc_idxminmax(
        # we need to attach back the dim name
        res.name = dim
    else:
+        indx.data = to_numpy(indx.data)


Not sure this is exactly what we want but got the cupy tests working.

Yeah, I would lean aginst this. It looks like cupy does not support indexing cupy arrays with cupy arrays?

Or is the issue that the cupy arrays are used to index xarray coordinates?

slevang · 2024-11-18T21:01:43Z

xarray/core/duck_array_ops.py

-)
+
+def sliding_window_view(array, window_shape, axis=None, **kwargs):
+    # TODO: some libraries (e.g. jax) don't have this, implement an alternative?


This is one of the biggest outstanding bummers of wrapping jax arrays. There is apparently openness to adding this as an API (even though it would not offer any performance benefit in XLA). But given this is way outside the API standard, whether it makes sense to implement a general version within xarray that doesn't rely on stride tricks.

You could implement a version using "summed area tables" (basically run a single accumulator and then compute differences between the window edges); or convolutions I guess.

I have something that works pretty well with this style of gather operation. But only in a jit context where XLA can work its magic. So I guess this is better left to the specific library to implement, or the user.

slevang · 2024-11-18T21:13:40Z

xarray/tests/test_duck_array_wrapping.py

+
+import xarray as xr
+
+# TODO: how to test these in CI?


note: I just noticed xarray-array-testing

lucascolley · 2024-11-18T21:44:44Z

see also data-apis/array-api#621 for the higher level discussion of nan reductions

lucascolley · 2024-11-18T21:50:14Z

xarray/core/array_api_compat.py

+def get_array_namespace(*values):
+    def _get_single_namespace(x):
+        if hasattr(x, "__array_namespace__"):
+            return x.__array_namespace__()
+        elif isinstance(x, array_type("cupy")):
+            # special case cupy for now
+            import cupy as cp
+
+            return cp
+        else:
+            return np


you probably want to wrap array-api-compat's array_namespace here, see https://github.com/scipy/scipy/blob/ec30b43e143ac0cb0e30129d4da0dbaa29e74c34/scipy/_lib/_array_api.py#L118-L152 for what we do in SciPy.

Happy to go that route, I did actually try that but array-api-compat doesn't handle a bunch of things we end up passing through this (scalars, index wrappers, etc) so it would require some careful prefiltering.

The only things this package effectively wraps that don't have __array_namespace__ are cupy, dask, and torch. I tried the torch wrapper but it doesn't pass our as_compatible_data check because the wrapper object itself doesn't have __array_namespace__ or __array_function__ 😕

Support for scalars was just merged half an hour ago! data-apis/array-api-compat#147

Looks like this would handle the array-like check well, but would require adding this as a core xarray dependency to use it in as_compatible_data. Not sure if there is any appetite for that.

In SciPy we vendor array-api-compat via a Git submodule. There's a little bit of build system bookkeeping needed, but otherwise it works well without introducing a dependency.

I played around with this some more.

Getting an object that is compliant from the perspective of array_api_compat.is_array_api_obj to work through the hierarchy of xarray objects requires basically swapping in a bunch of more restrictive hasattr(__array_namespace__) checks for this function. Not too bad.

I was able to get things to the point that xarray can wrap torch.Tensor, which is pretty cool. But the reality is that xarray relies on a lot of functionality beyond the Array API, so this just doesn't work in any practical sense. Similarly, swapping in the more restrictive cupy.array_api for the main cupy namespace causes all kinds of things to break.

It seems torch has had little movement on supporting the standard since 2021. From xarray's perspective, cupy implements everything we need, except for the actual __array_namespace__ attribute declaring compatibility, which seems to be planned for v14.

So while this compat module is nice in theory, I don't think it's very useful for xarray. cupy, jax, and sparse are good to go without it, and we only need a single special case for cupy to fetch its __array_namespace__.

slevang · 2024-11-20T15:47:08Z

I looked back at cupy-xarray. Now that this duck array stuff all works pretty well, I'm wondering how people feel about adding official DataArray/Dataset methods analogous to this, but in a generalized way. Something like:

def as_array_type(self, asarray: callable, **kwargs) -> Self:
# e.g. ds.to_array_type(jnp.asarray, device="gpu")

def is_array_type(self, array: type) -> bool:
# e.g. ds.is_array_type(cp.ndarray)

dcherian · 2024-11-20T17:51:57Z

methods analogous to this, but in a generalized way.

👍 I though I'd seen an array API method for converting between compliant array types, but I can't find it now

slevang · 2024-11-20T18:07:41Z

Compliant namespaces should now implement from_dlpack which is generally the recommended conversion protocol. So I suppose we could instead pass the namespace and hard code it to use that:

def to_namespace(self, xp: ModuleType, **kwargs) -> Self:
    xp.from_dlpack(self.data)
       
# e.g. ds.to_namespace(cp)

But this actually doesn't work for cupy, since they're quite stringent about implicit device transfers:

TypeError: CPU arrays cannot be directly imported to CuPy. Use `cupy.array(numpy.from_dlpack(input))` instead.

Also sparse doesn't have this at all, and it wouldn't be clear whether you want a COO, DOK, etc.

shoyer · 2024-11-20T20:59:01Z

xarray/core/computation.py

@@ -2174,9 +2174,13 @@ def _calc_idxminmax(
        # we need to attach back the dim name
        res.name = dim
    else:
+        indx.data = to_numpy(indx.data)


Yeah, I would lean aginst this. It looks like cupy does not support indexing cupy arrays with cupy arrays?

Or is the issue that the cupy arrays are used to index xarray coordinates?

shoyer · 2024-11-20T21:02:21Z

xarray/core/dataset.py

+                    if isinstance(v.data, array_type("cupy")):
+                        coord_data = duck_array_ops.get_array_namespace(v.data).asarray(
+                            coord_var.data
+                        )
+                    else:
+                        coord_data = coord_var.data


It looks like this could be encapsulated into an as_like_array() helper function?

I'd like to keep the explicit array namespace stuff out of most xarray functions.

shoyer · 2024-11-20T21:03:51Z

xarray/core/variable.py

+            if type(mask) is not type(data):
+                mask = duck_array_ops.get_array_namespace(data).asarray(mask)
+            data = duck_array_ops.where(
+                duck_array_ops.logical_not(mask), data, fill_value
+            )


this could also use as_like_array()

slevang · 2024-11-21T16:22:15Z

I think this is in pretty good shape now, except the question of whether to attempt any of this integration testing in CI. That could also be punted to xarray-array-testing (@keewis @TomNicholas)

slevang added 5 commits November 18, 2024 11:03

lots more duck array compat, plus tests

fd6b339

Merge branch 'main' into more-array-api-compat

893408c

merge sliding_window_view

f7866ce

namespaces constant

90037fe

revert dask allowed

5ba1a2f

slevang commented Nov 18, 2024

View reviewed changes

slevang mentioned this pull request Nov 18, 2024

Compatibility with the Array API standard #7848

Open

lucascolley reviewed Nov 18, 2024

View reviewed changes

slevang added 2 commits November 19, 2024 15:20

fix up some tests

6225ae3

backwards compat sparse mask

e2911c2

slevang marked this pull request as ready for review November 20, 2024 15:47

shoyer reviewed Nov 20, 2024

View reviewed changes

slevang added 11 commits November 20, 2024 22:13

add as_array methods

2ac37f9

to_like_array helper

1cc344b

Merge branch 'main' into more-array-api-compat

69080a5

only cast non-numpy

372439c

better idxminmax approach

0eef2cb

fix mypy

6739504

naming, add is_array_type

9e6d6f8

add public doc and whats new

e721011

update comments

1fe4131

add support for chunked arrays in as_array_type

205c199

Merge branch 'main' into more-array-api-compat

7752088

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improved duck array wrapping #9798

Improved duck array wrapping #9798

slevang commented Nov 18, 2024 •

edited

Loading

slevang Nov 18, 2024

slevang Nov 18, 2024

shoyer Nov 20, 2024

slevang Nov 18, 2024

dcherian Nov 18, 2024

slevang Nov 20, 2024

slevang Nov 18, 2024

lucascolley commented Nov 18, 2024

lucascolley Nov 18, 2024

slevang Nov 18, 2024

lucascolley Nov 18, 2024

slevang Nov 19, 2024

lucascolley Nov 19, 2024

slevang Nov 20, 2024

slevang commented Nov 20, 2024

dcherian commented Nov 20, 2024

slevang commented Nov 20, 2024

shoyer Nov 20, 2024

shoyer Nov 20, 2024

shoyer Nov 20, 2024

slevang commented Nov 21, 2024

Improved duck array wrapping #9798

Are you sure you want to change the base?

Improved duck array wrapping #9798

Conversation

slevang commented Nov 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lucascolley commented Nov 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slevang commented Nov 20, 2024

dcherian commented Nov 20, 2024

slevang commented Nov 20, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slevang commented Nov 21, 2024

slevang commented Nov 18, 2024 •

edited

Loading