Add __array_ufunc__ protocol to dask.array by mrocklin · Pull Request #2438 · dask/dask

mrocklin · 2017-06-07T21:01:22Z

This implements the __array_ufunc__ protocol for dask.array. Generally we just dispatch to the already-existing functions. I did this in haste and without properly reviewing the NEP/PR so I would not be surprised if I've missed something.

Example

In [1]: import dask.array as da, numpy as np

In [2]: x = np.arange(24).reshape((4, 6))

In [3]: np.sum(np.sin(x), axis=0)
Out[3]:
array([-1.56697566,  2.06850183,  3.80220828,  2.04018197, -1.59757823,
       -3.76653238])

In [4]: d = da.arange(24, chunks=(4,)).reshape((4, 6))

In [5]: np.sum(np.sin(d), axis=0)
Out[5]: dask.array<sum-aggregate, shape=(6,), dtype=float64, chunksize=(6,)>

In [6]: _.compute()
Out[6]:
array([-1.56697566,  2.06850183,  3.80220828,  2.04018197, -1.59757823,
       -3.76653238])

cc @shoyer @njsmith

mrocklin · 2017-06-07T21:34:21Z

Now also supporting the out= parameter. Interestingly this fails for cumulative reductions, which seem to be assuming numpy arrays.

In [1]: import dask.array as da

In [2]: import numpy as np

In [3]: x = da.ones((15, 15), chunks=(5, 5))

In [4]: np.add(x, 100, out=x)

In [5]: empty = da.zeros((15,), chunks=(5,))

In [6]: np.sum(x, axis=0, out=empty)

In [7]: empty.visualize('dask.png')
Out[7]: <IPython.core.display.Image object>

In [8]: empty.compute()
Out[8]: 
array([ 1515.,  1515.,  1515.,  1515.,  1515.,  1515.,  1515.,  1515.,
        1515.,  1515.,  1515.,  1515.,  1515.,  1515.,  1515.])

shoyer

Also, you probably realize this, but just to note that __array_ufunc__ is only used for NumPy ufunc method. It's not used for reductions like sum or cumsum, though those use an ad-hoc dispatch system based on calling methods.

shoyer · 2017-06-08T00:38:56Z

+        from . import ufunc
+        da_ufunc = getattr(ufunc, numpy_ufunc.__name__)
+        if method == '__call__':
+            return da_ufunc(*inputs, **kwargs)


I would suggest writing just elemwise(numpy_ufunc, *inputs, **kwargs), once you've verified numpy_ufunc.signature is None (which means that isn't a gufunc). This would allow you to handle arbitrary ufuncs without dask wrappers.

I've done this in most cases except those where the ufunc returns multiple outputs. In that case I'm reverting back to calling the dask.array version.

shoyer · 2017-06-08T00:40:55Z

+            return da_ufunc.outer(*inputs, **kwargs)
+        else:
+            from ..base import compute
+            inputs2 = compute(*inputs)


I'm not sure this is a good idea.

What would you recommend? I've changed this to returning NotImplemented for now.

shoyer · 2017-06-08T00:45:26Z

+    def __array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs):
+        import pdb; pdb.set_trace()
+        from . import ufunc
+        da_ufunc = getattr(ufunc, numpy_ufunc.__name__)


You need to return NotImplemented here for unrecognized types, e.g., something along these lines. Otherwise, dask will take precedence over everything else, even libraries like xarray.

(There are also some subtleties around output types and the out argument that I would recommend copying from that example, though possibly it would make sense to always return NotImplemented if out is set.)

We're now returning NotImplemented when we don't know how to handle the inputs.

TomAugspurger

Looks great. This will need a pretty big warning for users who (maybe unknowingly via a library) relied on numpy ufuncs evaluating dask arrays eagerly.

TomAugspurger · 2017-06-08T13:20:55Z

+    def __array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs):
+        out = kwargs.get('out', ())
+        for x in inputs + out:
+            if not isinstance(x, (np.ndarray, Number, Array)):


What's the desired behavior when interacting with, say xarray.DataArrays? I think let them decide how to handle the inputs (so NotImplemented is correct here).

Yes, that's exactly why we need this.

Maybe add the sparse array type that is allowed in chunks here?

TomAugspurger · 2017-06-08T13:25:43Z

+    return handle_out(out, result)
+
+
+def handle_out(out, result):


This maybe could be implemented as a decorator, when functions have multiple return points (though I see you changed the one case where that happened). If elemwise was the only one that needed to be changed, then a regular function seems preferable.

I'm somewhat cautious about using decorators for internal things like this. I think that we're capable of tracking down the few corner cases here and I would be concerned about reducing the explicit-ness of the code. This is pretty subjective though and just a personal preference.

mrocklin · 2017-06-08T13:51:32Z

This will need a pretty big warning for users who (maybe unknowingly via a library) relied on numpy ufuncs evaluating dask arrays eagerly

Perhaps we can squeeze a FutureWarning about this behavior into the upcoming release.

shoyer · 2017-06-08T15:54:23Z

+    def __array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs):
+        out = kwargs.get('out', ())
+        for x in inputs + out:
+            if not isinstance(x, (np.ndarray, Number, Array)):


Yes, that's exactly why we need this.

shoyer · 2017-06-08T15:56:16Z

+        if method == '__call__':
+            if numpy_ufunc.signature is not None:
+                return NotImplemented
+            if numpy_ufunc.__name__ in ('frexp', 'modf'):  # multi-output


You should check numpy_ufunc.nout, so you can at least raise for other ufuncs with multiple outputs.

Thanks, that's much cleaner

shoyer · 2017-06-08T15:57:04Z

+                return elemwise(numpy_ufunc, *inputs, **kwargs)
+        elif method == 'outer':
+            from . import ufunc
+            da_ufunc = getattr(ufunc, numpy_ufunc.__name__)


Slightly safer:

try: da_ufunc = getattr(ufunc, numpy_ufunc.__name) except AttributeError: return NotImplemented

shoyer · 2017-06-08T16:03:30Z

+    def __array_ufunc__(self, numpy_ufunc, method, *inputs, **kwargs):
+        out = kwargs.get('out', ())
+        for x in inputs + out:
+            if not isinstance(x, (np.ndarray, Number, Array)):


Maybe add the sparse array type that is allowed in chunks here?

mrocklin · 2017-06-08T16:08:33Z

This will need a pretty big warning for users who (maybe unknowingly via a library) relied on numpy ufuncs evaluating dask arrays eagerly

Perhaps we can squeeze a FutureWarning about this behavior into the upcoming release.

OK, this might be hard. The only thing that gets triggered on our end is the __array__ call. I don't think that we should warn generally on this method. I don't know how to test for "we are being called by a numpy ufunc".

mrocklin · 2017-06-08T16:26:08Z

If we can't meaningfully provide warnings then do we want this change in the upcoming release?

TomAugspurger · 2017-06-08T16:32:07Z

OK, this might be hard. The only thing that gets triggered on our end is the array call.

Correct me if I'm wrong, but could the next release include a simple definition for __array_ufunc__ that just emits a FutureWarning and then return NotImplemented so that Numpy takes over again? Or would that have poor interaction with other inputs that implement __array_ufunc__?

mrocklin · 2017-06-08T16:33:15Z

That sounds good. I was trying to solve this for older versions of numpy, but that's probably not the right way to think about this.

shoyer · 2017-06-08T16:33:27Z

I think the way to provide warnings would be to implement a version of __array_ufunc__ that calls compute on each dask array argument and then calls the ufunc again.

But I'm not entirely sure this is worth the trouble -- possibly better to save this for a breaking release.

mrocklin · 2017-06-08T17:17:19Z

We could easily make the next release such a breaking release. I'm fine either way. I think that having a release with a warning is probably the mature thing to do. However I also don't know of anyone who intentionally uses numpy ufuncs on dask.arrays and there would be some community value in getting a project out that used the new NumPy 1.13.0 features.

mrocklin · 2017-06-08T19:39:59Z

OK, I'm inclined to use this as an opportunity to break a couple of other things in dask/distributed and issue a larger release.

mrocklin · 2017-06-09T00:55:57Z

Looks like we're no longer raising a NotImplementedError when using the matmul operator. I'm not quite sure why this is.

    def test_matmul():
        x = np.random.random((5, 5))
        y = np.random.random((5, 2))
        a = from_array(x, chunks=(1, 5))
        b = from_array(y, chunks=(5, 1))
        assert_eq(operator.matmul(a, b), a.dot(b))
        assert_eq(operator.matmul(a, b), operator.matmul(x, y))
        assert_eq(operator.matmul(a, y), operator.matmul(x, b))
        list_vec = list(range(1, 6))
        assert_eq(operator.matmul(list_vec, b), operator.matmul(list_vec, y))
        assert_eq(operator.matmul(x, list_vec), operator.matmul(a, list_vec))
        z = np.random.random((5, 5, 5))
        with pytest.raises(NotImplementedError):
            operator.matmul(a, z)
        c = from_array(z, chunks=(1, 5, 1))
        with pytest.raises(NotImplementedError):
>           operator.matmul(x, c)
E           Failed: DID NOT RAISE <class 'NotImplementedError'>

shoyer · 2017-06-09T01:02:42Z

This looks like numpy/numpy#9028

Out of curiosity, what does the result of operator.matmul(x, c) look like?

mrocklin · 2017-06-09T01:03:34Z

A numpy array. It doesn't appear to go through the __array_ufunc__ code path.

Example ------- ```python In [1]: import dask.array as da, numpy as np In [2]: x = np.arange(24).reshape((4, 6)) In [3]: np.sum(np.sin(x), axis=0) Out[3]: array([-1.56697566, 2.06850183, 3.80220828, 2.04018197, -1.59757823, -3.76653238]) In [4]: d = da.arange(24, chunks=(4,)).reshape((4, 6)) In [5]: np.sum(np.sin(d), axis=0) Out[5]: dask.array<sum-aggregate, shape=(6,), dtype=float64, chunksize=(6,)> In [6]: _.compute() Out[6]: array([-1.56697566, 2.06850183, 3.80220828, 2.04018197, -1.59757823, -3.76653238]) ```

NumPy 1.13.0 is not yet compatible with pyarrow or pandas 0.20.1

mrocklin · 2017-06-10T15:20:30Z

OK, in order to add numpy 1.13.0 to the testing matrix I had to move around a couple things because some packages, notably on conda-forge, aren't yet compatible with the new version. In particular I made the following changes:

Installing PyArrow from PyPI
Avoid the new version of pytest for now, pinning pytest<=3.1.1
Install graphviz from conda-forge

These latter two aren't necessarily related, but seem like decent ideas regardless.

mrocklin · 2017-06-10T15:21:09Z

This now passes tests. I'm inclined to merge and release. @TomAugspurger are you comfortable with the lack of a deprecation cycle here if we bump the version by 0.X.0?

TomAugspurger · 2017-06-10T16:28:08Z

Yep, going forward seems like the right thing. My initial request for a warning was a warning in the release notes.

…

On Jun 10, 2017, at 10:21 AM, Matthew Rocklin ***@***.***> wrote: This now passes tests. I'm inclined to merge and release. @TomAugspurger are you comfortable with the lack of a deprecation cycle here if we bump the version by 0.X.0? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub, or mute the thread.

mrocklin · 2017-06-10T16:37:38Z

I plan to merge this in a few hours if there are no further comments

jcrist · 2017-06-10T17:59:55Z

 conda install -q -c conda-forge \
    distributed \
    cloudpickle \
+    graphviz \


~~Is this the c library? Why this change?~~

Edit: nvm, just saw your comment.

I think that your question remains valid. I was getting png failures from graphviz. I suspect that the defaults channel build is less-than-optimal again. I didn't want to deal with it just now though so punted to use conda-forge for the time being. I'll report upstream after the current build-storm passes if the problem persists.

mrocklin mentioned this pull request Jun 7, 2017

ENH: Add __array_ufunc__ numpy/numpy#8247

Merged

mrocklin requested a review from TomAugspurger June 7, 2017 21:54

mrocklin force-pushed the array_ufunc branch from e203498 to bc5e2c6 Compare June 7, 2017 23:58

shoyer reviewed Jun 8, 2017

View reviewed changes

mrocklin force-pushed the array_ufunc branch from 496008c to ed78009 Compare June 8, 2017 13:05

TomAugspurger reviewed Jun 8, 2017

View reviewed changes

shoyer reviewed Jun 8, 2017

View reviewed changes

shoyer mentioned this pull request Jun 9, 2017

BUG: no way to override matmul/@ if __array_ufunc__ is set numpy/numpy#9028

Closed

mrocklin added 10 commits June 8, 2017 22:10

Support out parameters

1dbc697

test non-supported ufunc methods

c1373c3

fix cumreductions

d5cde0c

Use elemwise in __array_ufunc__ and return NotImplemented

16ba7fc

Test out= paramter with numpy array

bbf1348

cleanup tests for numpy < 1.13.0

f5c9879

generalize special cases in __array_ufunc__

7b32bc7

flake8

c3f3c5e

skip out= test under old numpy

2f60221

mrocklin added 4 commits June 8, 2017 22:10

skip fix test

7cc0687

use numpy=1.13.0 in travis.yml

b6c0156

Update travis.yml libraries

fdac908

NumPy 1.13.0 is not yet compatible with pyarrow or pandas 0.20.1

pull out matmul raise issue and skip

17d3033

mrocklin force-pushed the array_ufunc branch from 0bb99ba to 17d3033 Compare June 9, 2017 02:37

mrocklin added 5 commits June 9, 2017 17:15

remove comment within travis install

63ee9eb

update travis.yml

6235022

add graphviz from conda-forge pyarrow from pip

b1989d3

pin pytest

6525f30

pin pytest below 3.1.1

0c566c1

jcrist reviewed Jun 10, 2017

View reviewed changes

mrocklin merged commit 72a5f51 into dask:master Jun 11, 2017

mrocklin deleted the array_ufunc branch June 11, 2017 13:08

sinhrks added this to the 0.15.0 milestone Aug 30, 2017

Conversation

mrocklin commented Jun 7, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Example

Uh oh!

mrocklin commented Jun 7, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shoyer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

TomAugspurger commented Jun 8, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

shoyer commented Jun 8, 2017

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

mrocklin commented Jun 8, 2017

Uh oh!

mrocklin commented Jun 9, 2017

Uh oh!

shoyer commented Jun 9, 2017

Uh oh!

mrocklin commented Jun 9, 2017

Uh oh!

mrocklin commented Jun 10, 2017

Uh oh!

mrocklin commented Jun 10, 2017

Uh oh!

TomAugspurger commented Jun 10, 2017 via email

Uh oh!

mrocklin commented Jun 10, 2017

Uh oh!

jcrist Jun 10, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

mrocklin commented Jun 7, 2017 •

edited

Loading

mrocklin commented Jun 7, 2017 •

edited

Loading

TomAugspurger commented Jun 8, 2017 •

edited

Loading

jcrist Jun 10, 2017 •

edited

Loading