Support for NEP-35 #6738

pentschev · 2020-10-15T12:56:09Z

This PR adds support for NEP-35. To describe the NEP briefly, it proposes the inclusion of a new keyword argument like to NumPy array creation functions, such as np.array, allowing NumPy to dispatch the array creation downstream to other libraries, such as CuPy and Dask itself. To demonstrate this, consider the following example:

>>> import numpy as np, cupy as cp, dask.array as da
>>>
>>> cp_ref = cp.array(0)
>>> da_ref = da.array(0)
>>>
>>> type(np.array(da.arange(5)))
<class 'numpy.ndarray'>
>>> type(np.array(da.arange(5), like=da_ref))
<class 'dask.array.core.Array'>
>>>
>>> type(np.array(cp.arange(5)))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "cupy/core/core.pyx", line 1245, in cupy.core.core.ndarray.__array__
TypeError: Implicit conversion to a NumPy array is not allowed. Please use `.get()` to construct a NumPy array explicitly.
>>> type(np.array(cp.arange(5), like=cp_ref))
<class 'cupy.core.core.ndarray'>

NEP-35 is still experimental and requires NumPy 1.20.0+.

The changes here attempt to cover a large portion of Dask, adding some new tests that wouldn't pass before. This change is mostly beneficial for downstream libraries such as CuPy and Sparse, that implement the __array_function__ protocol, allowing Dask to continue working with, for example, Dask arrays backed by CuPy in functions such as pad(..., mode="constant"), which wasn't possible before because of the call to np.asarray that wasn't able to create a new array in a type other than NumPy.

Things that are not yet supported include (but not limited to):

da.random.choice: it's a bit of complex function which will take some time to get right, but doesn't seem to have any technical limitations;
linalg functions: requires solve_triangular which is a SciPy function and can't be dispatched with __array_function__, thus will probably require special casing;

cc @quasiben @jakirkham @mrocklin @seberg @shoyer @rgommers @hameerabbasi @leofang

quasiben · 2021-02-11T13:40:13Z

Yay tests pass 🎉 Thank you for continuing to push here @pentschev!

Is it fair to characterize most of the changes here as moving the type checks from isinstance(arr, list, np.ndarray) to is_arraylike(arr) and updating safety function calls to use the like= arg:

asarray_safe(c, like=meta_from_array(array), dtype=result.dtype)

I think this is a large enough change that we should ask for an external reviewer if folks have time. Ideally we could merge in tomorrow but that might not be enough lead time. If we don't hear back I think planning for early next week makes sense.

pentschev · 2021-02-11T15:10:55Z

Is it fair to characterize most of the changes here as moving the type checks from isinstance(arr, list, np.ndarray) to is_arraylike(arr) and updating safety function calls to use the like= arg:

asarray_safe(c, like=meta_from_array(array), dtype=result.dtype)

For the first part, no. Take for example https://github.com/dask/dask/pull/6738/files#diff-cc8bb1c7ea02a9582c31ac1580f3ab0674c35e0b6af060513ba9932ceb07a600R227, that's a place where support for non-arraylike is still necessary, thus that has to be verified on a per-case basis.

For the second part, yes. The intent is that *array_safe should be backwards compatible too, much like what has been done with empty_like_safe and its counterparts, rather than having devs writing np.array*(arr, like=some_arr) and then checking whether it succeeds (NumPy 1.20+) or fails for older NumPy. I'm not currently aware of any cases for which *array_safe doesn't reproduce a backwards-compatible behavior, but if there are any, I'm confident we can fix that such that using *array_safe should be preferred over np.array*.

hameerabbasi · 2021-02-11T15:13:21Z

dask/array/tests/test_cupy.py

+    assert_eq(result, np.array([1], dtype=d.dtype))
+
+
+@pytest.mark.skipif(np.__version__ < "1.20", reason="NEP-35 is not available")


I would move this into a global variable, perhaps in a settings-like module. That way it's easy to change the test.

Specifically, create a variable named NEP35_AVAILABLE.

Good point, this exists already in

dask/dask/array/numpy_compat.py

Line 11 in ac2e582

_numpy_120 = LooseVersion(np.__version__) >= "1.20.0"

which should suffice.

Done in bae0dd6

I would suggest testing the protocol rather than the version. Something like:

def _test_nep35(): class MyClass: def __array_function__(self, func, args, kwargs): return True try: return np.asarray([], like=MyClass()) except TypeError: return False NEP35_AVAILABLE = _test_nep35()

I would personally prefer they fail to force us to update the tests should anything change, rather than having tests getting skipped. Happy to change my mindset should there be compelling arguments in favor of skipping. 🙂

quasiben · 2021-02-16T21:17:28Z

I think @crusaderky may take a look at this -- if not, I'll plan to merge tomorrow morning

jakirkham · 2021-02-17T05:03:58Z

I’ve been meaning to take a look at this as well. Will carve out some time tomorrow

quasiben · 2021-02-18T22:34:16Z

I'm going to merge now. If there are additional concerns I'm sure @pentschev would be happy to follow up with PRs/comments/etc.

seberg · 2021-02-18T22:54:37Z

Small heads-up: we missed some small memory leaks in np.empty and np.zeros when like= is used, those should be gone in 1.20.2, but in case someone notices...

jakirkham · 2021-02-18T22:55:24Z

Sorry for the delay. Skimmed over the changes. These seem reasonable. No additional comments

jakirkham · 2021-02-18T23:44:50Z

Would it makes sense to do a blogpost on this? Like a short follow on/update of the NEP 18 one?

pentschev · 2021-02-19T14:50:35Z

Thanks @seberg , is numpy/numpy@e907085 the fix you mentioned?

seberg · 2021-02-19T14:53:52Z

Yeah, it just leaks "shape" information, so thats a couple a few dozen bytes each call (don't worry about the dtype, in the vast majority of cases those are singletons anyway).

pentschev · 2021-02-19T15:09:40Z

Thanks @seberg for catching and fixing that, it's my fault to have missed it! 🙁

pentschev · 2021-02-19T15:09:51Z

Would it makes sense to do a blogpost on this? Like a short follow on/update of the NEP 18 one?

I'm not sure, possibly? This change is generally more transparent to users than NEP-18 was, so I'm not exactly sure how to approach it in a manner that anybody who's not involved in library development would be compelled to read. No objections in discussing it, I just don't see a good way to write it right now.

quasiben and others added 30 commits August 31, 2020 12:58

use nep35 in slicing

3f5fdb6

Merge branch 'master' of github.com:dask/dask into nep-35-test

c8f406e

add compress test

8a2c314

Merge remote-tracking branch 'upstream/master' into nep-35-test

e89c87b

Add asanyarray_safe function supporting like= kwarg

c02d94a

Add CuPy tests for asanyarray_safe

e3862b6

Use asanyarray_safe in slicing functions

f12b1c1

Add meta support to broadcast_to

cc19379

Support like= in pad with constant values

91e5256

Add CuPy tests for pad

9404577

Added asarray_safe, move implementation to private function

75bc6d9

Support like= in pad with linear_ramp

f82ee95

Support for like= in digitize

29d2a0c

Add CuPy test for digitize

a1c8d13

Support for like= in take

d4fe2fa

Support for meta= in _vindex

86e24a3

Add CuPy test for _vindex

c29fbe5

Support for like= in percentile

7820305

Merge remote-tracking branch 'upstream/master' into nep-35-test

1811732

Add array_safe

4e1a35a

Add CuPy tests for percentiles

f07df2e

Fixes for like= in slicing

c78db6e

Support for like= in view

f9ef85c

Add CuPy tests for view

43a41cc

Additional like= support in slicing

a20dded

Support for like= in getter, load_store_chunk

da6ac0d

Add CuPy tests for getter, store

7ae85b2

Fix raise error in CuPy test for asanyarray

276c52b

Merge remote-tracking branch 'upstream/master' into nep-35-test

05031c7

Undo accidental changes to test_sparse_hstack_vstack_csr

354651a

Fix black formatting

0e5cbc7

jakirkham mentioned this pull request Feb 11, 2021

GPU Support dask/dask-image#133

Open

hameerabbasi reviewed Feb 11, 2021

View reviewed changes

Update CuPy tests skipif condition with not _numpy_120

bae0dd6

pentschev mentioned this pull request Feb 15, 2021

dask.pad does not work with cupy arrays #6117

Closed

pentschev mentioned this pull request Feb 16, 2021

Array._meta missing functionality tracker #4888

Closed

20 tasks

quasiben merged commit 5948ec3 into dask:master Feb 18, 2021

pentschev mentioned this pull request Feb 22, 2021

DOC: update NEP status for accepted/finished NEPs numpy/numpy#18458

Merged

jrbourbeau mentioned this pull request Feb 23, 2021

Unstable dask causes unexpected compute results from xarray #7263

Closed

mathause mentioned this pull request Feb 23, 2021

Upstream CI failing silently pydata/xarray#4945

Closed

pentschev deleted the nep-35-support branch February 23, 2021 20:31

This was referenced Mar 5, 2021

Failing CuPy tests #7324

Closed

NEP-35 duck array update #7321

Merged

This was referenced Mar 8, 2021

Documentation build fails with latest Dask release (2021.3.0) sgkit-dev/sgkit#482

Closed

Pin to Dask 2021.02.0 while incompatibility with later versions is investigated sgkit-dev/sgkit#483

Merged

ericpre mentioned this pull request Mar 17, 2021

Support CUDA GPU accelerated template matching pyxem/pyxem#743

Merged

3 tasks

quasiben mentioned this pull request Mar 24, 2021

GPU CI dask/community#138

Closed

pentschev mentioned this pull request May 12, 2021

Simple code is reaching a maximum recursion depth exception #7644

Closed

		assert_eq(result, np.array([1], dtype=d.dtype))


		@pytest.mark.skipif(np.__version__ < "1.20", reason="NEP-35 is not available")

Uh oh!

Support for NEP-35 #6738

Support for NEP-35 #6738

Uh oh!

Conversation

pentschev commented Oct 15, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

quasiben commented Feb 11, 2021

Uh oh!

pentschev commented Feb 11, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

hameerabbasi Feb 11, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hameerabbasi Feb 11, 2021

Choose a reason for hiding this comment

Uh oh!

pentschev Feb 11, 2021

Choose a reason for hiding this comment

Uh oh!

pentschev Feb 11, 2021

Choose a reason for hiding this comment

Uh oh!

hameerabbasi Feb 11, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pentschev Feb 11, 2021

Choose a reason for hiding this comment

Uh oh!

quasiben commented Feb 16, 2021

Uh oh!

jakirkham commented Feb 17, 2021

Uh oh!

quasiben commented Feb 18, 2021

Uh oh!

seberg commented Feb 18, 2021

Uh oh!

jakirkham commented Feb 18, 2021

Uh oh!

jakirkham commented Feb 18, 2021

Uh oh!

pentschev commented Feb 19, 2021

Uh oh!

seberg commented Feb 19, 2021

Uh oh!

pentschev commented Feb 19, 2021

Uh oh!

pentschev commented Feb 19, 2021

Uh oh!

Uh oh!

pentschev commented Oct 15, 2020 •

edited

Loading

pentschev commented Feb 11, 2021 •

edited

Loading

hameerabbasi Feb 11, 2021 •

edited

Loading

hameerabbasi Feb 11, 2021 •

edited

Loading