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?

xarray/core/computation.py

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

we do install pint in CI if that helps.

Cupy requries a GPU, but we could add jax into Xarray's CI suite in GitHub actions, at least.

the only thing we'll want to avoid is testing compat with too many array types in xarray's CI. I guess I'll need to make progress on xarray-array-testing to avoid that.

jax would be pretty easy to add. The CPU jaxlib package is only like 90MB. And it's a nice test case since it very closely matches the numpy API, but only implements __array_namespace__, and not __array_function__ or __array_ufunc__.

I agree, we should add JAX CPU. See https://github.com/pydata/xarray/tree/main/ci/requirements for the config files. Feel free to do this in this PR or a follow-up.

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.

xarray/core/computation.py

xarray/core/dataset.py

xarray/core/variable.py

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)

shoyer · 2024-11-22T18:14:43Z

The new is_array_type and as_array_types methods seem worthy of broader discussion. I would leave them out of this PR for now, so they don't hold up merging it.

slevang · 2024-11-22T18:28:28Z

Sure I'll split that off. Agree the naming and API could use some discussion.

dcherian · 2024-11-22T21:08:27Z

Basically everything works except the interp/missing methods which have a lot of specialized code.

interp will need scipy to support array-api objects which I think they're working on.

Also a few odds and ends like polyfit and rank.

yeah funny that the linalg extension does not include lstsq.

More generally, this is great work @slevang. thanks!

lucascolley · 2024-11-22T21:12:53Z

yeah funny that the linalg extension does not include lstsq.

From 2021:

Initial discussion on linalg.lstsq API design can be found in data-apis/array-api#227. There, it was held that the current API design posed undue burden on specification implementers due to API complexity (returning solutions, ranks, residuals, and singular values) and violated various linear algebra API design principles (particularly orthogonality).

Based on ensuing discussion in consortium meetings, consensus was found in simply removing the linalg.lstsq API altogether. This API finds low usage among downstream libraries (see data) and has several variants due to specialized use cases (see SciPy).

Should this or similar functionality be found to be desirable in the future, specification can happen in a future revision of the standard.

You aren't the first to hit it, reading that issue. Perhaps you could raise it again :)

dcherian · 2024-11-22T21:14:30Z

aha thanks @lucascolley . Indeed we should rewrite to use solve probably.

slevang · 2024-11-22T21:26:31Z

fwiw, jax and cupy both have linalg.lstsq, although it sounds like solve would be the right way to go.

I just didn't tackle this function because it is defined on Dataset, and has probably 15 calls to different numpy functions, followed by a loop over variables. So it seemed tricky to handle in case we have data vars with mixed namespaces or other edge cases.

shoyer · 2024-11-22T21:36:16Z

xarray/core/duck_array_ops.py

@@ -172,7 +166,8 @@ def isnull(data):
        )
    ):
        # these types cannot represent missing values
-        return full_like(data, dtype=bool, fill_value=False)
+        dtype = xp.bool_ if hasattr(xp, "bool_") else xp.bool


I believe xp.bool_ is only needed for compatibility with NumPy 1.x releases. Can we make that clearer here, e.g., by using np.bool_ and adding a comment about NumPy versions?

Will do, but also cupy only has bool_ at the moment. But If I passed built-in bool to this I think I got a cupy error. It was a weird range from like numpy 1.2x until 2 that they deprecated bool and only had bool_, then brought it back in 2.0 🤷

shoyer · 2024-11-22T21:39:17Z

xarray/tests/test_duck_array_wrapping.py

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


Cupy requries a GPU, but we could add jax into Xarray's CI suite in GitHub actions, at least.

slevang · 2024-11-23T18:04:51Z

xarray/tests/test_duck_array_wrapping.py

+        },
+        "xfails": {"quantile": "no nanquantile"},
+    },
+    "pint": {


Added pint and sparse, and more consolidated xfail lists. Some of these are probably desired behavior, like all the methods that return bool or int not returning a pint.Quantity, but it was easiest to just xfail them.

slevang · 2024-11-23T18:06:59Z

ci/requirements/environment.yml

@@ -49,3 +49,5 @@ dependencies:
  - toolz
  - typing_extensions
  - zarr
+  - pip:
+      - jax # no way to get cpu-only jaxlib from conda if gpu is present


Trying to install via conda on my local machine with a GPU automatically gave all the CUDA libs, which seems like unnecessary bloat for dev environments. If anyone knows a way around this, happy to move it back to conda. I also got a failure on windows, I think for trying to pip install anything, so left it off of there.

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

xarray/core/computation.py Outdated Show resolved Hide resolved

xarray/core/dataset.py Outdated Show resolved Hide resolved

xarray/core/variable.py Outdated Show resolved Hide resolved

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

slevang added 3 commits November 22, 2024 14:19

revert array_type methods

c8d4e5e

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

e67a819

fix up whats new

f306768

dcherian approved these changes Nov 22, 2024

View reviewed changes

shoyer approved these changes Nov 22, 2024

View reviewed changes

slevang added 6 commits November 22, 2024 17:22

comment about bool_

18ebdcd

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

f51e3fb

add jax to complete ci envs

121af9e

add pint and sparse to tests

472ae7e

remove from windows

5aa4a39

mypy, xfail one more sparse

390df6f

slevang commented Nov 23, 2024

View reviewed changes

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

dcherian Nov 18, 2024

slevang Nov 20, 2024

slevang Nov 18, 2024

dcherian Nov 22, 2024

shoyer Nov 22, 2024

keewis Nov 22, 2024 •

edited

Loading

slevang Nov 22, 2024

shoyer Nov 22, 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

slevang commented Nov 21, 2024

shoyer commented Nov 22, 2024 •

edited

Loading

slevang commented Nov 22, 2024

dcherian commented Nov 22, 2024

lucascolley commented Nov 22, 2024

dcherian commented Nov 22, 2024

slevang commented Nov 22, 2024

shoyer Nov 22, 2024

slevang Nov 22, 2024

shoyer Nov 22, 2024

slevang Nov 23, 2024

slevang Nov 23, 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

keewis Nov 22, 2024 • edited Loading

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

slevang commented Nov 21, 2024

shoyer commented Nov 22, 2024 • edited Loading

slevang commented Nov 22, 2024

dcherian commented Nov 22, 2024

lucascolley commented Nov 22, 2024

dcherian commented Nov 22, 2024

slevang commented Nov 22, 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

slevang commented Nov 18, 2024 •

edited

Loading

keewis Nov 22, 2024 •

edited

Loading

shoyer commented Nov 22, 2024 •

edited

Loading