Duck array compatibility meeting

[TomNicholas]

Proposal: hold a high-level inter-library meeting to sort out roadblocks in the duck-array wrapping efforts.

Whilst trying to get dask, pint and xarray all working nicely together, I couldn't help but notice there are important issues which conclude with a shared sentiment that "we just need to make a decision as to what wraps what" but since then have had essentially no codified consensus, and hence no progress for the past year. Multiply-nested duck-array wrapping is complicated and involves a lot of separate libraries (as this graph of potential wrappings shows), but could be an amazingly powerful feature!

I suggest that as asynchronous discussion hasn't moved this forward, we should instead hold a (hopefully one-off) meeting to make these high-level design decisions.

I'm happy to arrange the meeting, but for this to work we ideally need attendees who understand the issues from the perspective of each of the main libraries involved - some suggestions:

• xarray (@shoyer and @keewis)
• dask (@mrocklin?)
• pint (@jthielen)
• cupy? (@jacobtomlinson?)
• sparse? (@crusaderky?)
• pytorch?? (@rgommers??)

Possible Agenda (please suggest additions!):

• Which libraries should wrap which other libraries
• Repo/NEP/etc. for standardizing wrapping order and other future decisions
• Outstanding issues to tackle first

Background reading

• Basic idea of the numpy dispatch mechanism explained in a blog post
• @jthielen 's excellent overview comment, with links to relevant NEP's
• Pint's technical commentary on array type support

Some related issues (there are many more - please add)

• UserWarning when wrapping pint & dask arrays together #5559
• Consistent Handling of Type Casting Hierarchy #3950
• Design: nested _meta dask/dask#5329
• reprs of nested duck arrays dask/dask#6637
• calling methods on wrapped duck arrays dask/dask#6636
• construction of and interactions between different types of nested duck arrays dask/dask#6635

[keewis]

cc @hameerabbasi for sparse, and @amcnicho might also be interested

[mrocklin]

I would be happy to attend and look forward to what I'm sure will be a vigorous discussion :) Thank you for providing convenient links to reading materials ahead of time.

As a warning, my responsiveness to github comments these days is not what it used to be. If I miss something here then please forgive me.

[hameerabbasi]

I'd also be happy to attend. Keep in mind I'm in the CET timezone.

[rgommers]

I'm happy to join, seems interesting. And yes, I can say something about PyTorch. There probably isn't much to say though - PyTorch is unlikely to adopt __array_function__ at this point, just like JAX. And it doesn't seem critical for this hierarchy anyway - the fundamental array objects (PyTorch/CuPy/NumPy/Sparse/JAX arrays or tensors) do not have or need a class hierarchy, they are all at the bottom and should not be mixed.

The key thing here seems to be Dask <-> Xarray <-> Pint, unless I'm missing something?

[jacobtomlinson]

Happy to attend. It might also be useful to have @pentschev involved too.

[TomNicholas]

Thanks to everyone ho has expressed interest already! I'll give some more time for responses before starting to think about potential meeting times.

There probably isn't much to say though - PyTorch is unlikely to adopt array_function at this point, just like JAX.

Interesting - could you say a bit more? Looking at these two issues, it seemed more like the question was simply on hold until someone who wanted it badly enough came along?

jax-ml/jax#1565

pytorch/pytorch#22402

(I'm interested because I'm in the set of people who have an ambition to use these libraries (wrapped in xarray) at some point further down the line, but that might remain just an ambition 😅)

The key thing here seems to be Dask <-> Xarray <-> Pint, unless I'm missing something?

You're right, that was the primary motivation for this issue, but it also seemed like a good opportunity to get as many libraries talking (and their relationships codified) as possible, especially if one of the outcomes ends up being a type casting hierarchy resource/NEP.

[jthielen]

Count me in for the meeting!

---

Here are a few suggestions about possible topics to add to the agenda (based on linked issues/discussions), if we can fit it all in:

• Canonical/minimal API of a "duck array" and how to detect it (though may be superseded by NEPs 30 and 47 among others)
• Consistency of type deferral (i.e., between construction, binary ops, __array_ufunc__, __array_function__, and array modules...for example, these are uniform in Pint, but construction and array module functions are deliberately different from the others for Dask arrays)
• API for inter-type casting and changing what types are used in a nested array (e.g. sparse and other duck array issues #3245 and Add to_numpy() and as_numpy() methods #5568)
• How to handle unknown duck arrays
• Nested array reprs (both short and full)
• Best practices for "carrying through" operations belonging to wrapped types (i.e., doing Dask-related things to a Pint Quantity or xarray DataArray that contains a Dask array), even if multiple layers deep

---

Also, tagging a few other array type libraries and maintainers/contributors who may be interested (please ping the relevant folks if you know them):

• unyt (@ngoldbaum)
• astropy.units (??)
• NumPy masked arrays (??)
• scipp (@SimonHeybrock, xref NEP 18, physical units, uncertainties, and the scipp library? #3509)

(interesting side note is the first three of these are all ndarray subclasses right now...perhaps discussing the interplay between array subclassing and wrapping is in order too?)

[benbovy]

• Awkward? (@jpivarski)

[TomNicholas]

NumPy masked arrays (??)

Tagging @greglucas because of his work on numpy/numpy#16022 (comment) - at least I think a very common desired use case of multi-nested arrays would be xarray(pint(dask(np.masked)).

[pentschev]

I'm also happy to join the meeting. Thanks @TomNicholas for the initiative here and @jacobtomlinson for tagging me.

[greglucas]

Happy to hop on a call for this as well, thanks for organizing all of this @TomNicholas!

[jpivarski]

• Awkward? (@jpivarski)

I'm interested. Let us know when the time will be or if there's a poll for picking a time. Thanks!

[rgommers]

Interesting - could you say a bit more? Looking at these two issues, it seemed more like the question was simply on hold until someone who wanted it badly enough came along?

There is a significant backwards compatibility break when a library adds __array_ufunc__ and __array_function__. JAX maintainers were not comfortable with that. @shoyer wrote https://numpy.org/neps/nep-0037-array-module.html as a follow-up largely because of that. That NEP is effectively superceded by the array API standard (https://data-apis.org/array-api/latest/ and NEP 47). PyTorch has decided to adopt that and implementation is well underway. Experimental support in NumPy will land next week (complete except for linalg). CuPy and MXNet plan to follow that NumPy implementation. JAX and Dask not yet decided I believe, but likely to follow NumPy as well.

Canonical/minimal API of a "duck array" and how to detect it (though may be superseded by NEPs 30 and 47 among others)

This is basically what the array API standard provides (most functions follow NumPy, with deviations mostly where other libraries were already deviating because they could implement something on GPU, with a JIT, or with a non-strided memory model). __array_function__ has worked quite well for CuPy and Dask, because they follow the NumPy API almost to the letter, with only a couple of exceptions (e.g. 0-D array instead of array scalars in CuPy). JAX, PyTorch and MXNet all deviate much more, and since the NumPy API is not very well-defined (there's 1500+ public objects plus more semi-public ones), you'd have no guarantees about what works and what doesn't.

That said, __array_ufunc__ and __array_function__ aren't going anywhere. The RAPIDS ecosystem is invested in it and I believe largely happy with it. So adding Xarray and Pint to the mix sounds potentially interesting.