Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add rollaxis and moveaxis #4822

Merged
merged 11 commits into from Aug 12, 2019

Conversation

@TAdeJong
Copy link
Contributor

commented May 19, 2019

Attempts to fix #2559 by aliasing the numpy functions, which depend on the transpose method.
Although #2955 existed as a WIP, I moved forward with a new PR with the easier aliasing method.

  • Tests added / passed
  • Passes flake8 dask
@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented May 19, 2019

Travis fails on Python3.7 only with a recursion error that I do not entirely understand or reproduce locally...

Can the assert isinstance() in combination with the getattr be a problem?

@jrbourbeau
Copy link
Member

left a comment

Thanks for this PR @TAdeJong!

I haven't had a chance for a thorough review yet, but the Python 3.7 errors on Travis seem to originate from the new NumPy __array_function__ protocol. The protocol is currently experimental and turned off by default, but it can be enabled by setting the environment variable NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1 (which is done in the 3.7 Travis build). The implementation of da.moveaxis here calls np.moveaxis which, when the __array_function__ protocol is enabled, dispatches back to the da.moveaxis function and that's where the RecursionError is coming from.

cc @jakirkham

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented May 19, 2019

Thanks for the quick explanation, that makes sense now and then a straight up alias moveaxis=np.moveaxis would make sense as @mrocklin suggested.
One thing I am not sure about: would just plainly in the place where the def is now be the right place to do such an alias? I will try to adapt this tomorrow. (Getting late here)

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

Thanks for picking this up, @TAdeJong. Agree that trying a simple alias is a good idea. Though we may still run into issues with __array_function__. That said, it is something easy to try and if it doesn't work we can give it more thought.

TAdeJong added 2 commits May 20, 2019
Change to simple alias
Less code, hopefully less clashes with `__array_function__` protocol
(seems to work locally?)
@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented May 20, 2019

OK, so a simple alias at least prevents infinite recursion, but in the case of NUMPY_EXPERIMENTAL_ARRAY_FUNCTION='1', we have a
TypeError: no implementation found for 'numpy.moveaxis' on types that implement __array_function__: [<class 'dask.array.core.Array'>]
Seems the array_function actually makes things a bit worse. I have not wrapped my head around NEP 18 entirely yet, anyone any suggestions?

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

Thanks for trying this out @TAdeJong. Yeah this is what I was alluding to above. Will take a look and see what options we have to move forward. Though others should feel free to jump in if they have thoughts as well.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

The __array_function__ dispatches NumPy functions to other APIs based on the types of the arrays. In the case here, the array is a Dask array, so it will dispatch to Dask. Since you were wrapping around a NumPy function, NumPy will dispatch to Dask, which will in turn call the NumPy function, and you get a recursion. If you only alias it, then there's no Dask implementation, as you've mentioned in your last comment.

I think the only solution for now would be to dispatch work through some mechanism such as da.map_blocks. I'm not too familiar with such Dask mechanisms, so maybe that isn't the most appropriate and there may be a better one for this particular case.

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

Thanks @pentschev. Your description of the problem sounds correct. Though we might have a better option for implementation.

After taking a quick look, it seems all NumPy functions as of NEP 18 have a __skip_array_function__ ( numpy/numpy#13389 ) method. I think if we try to alias that it should work to bypass NEP 18. However we will need to handle the case that an older version of NumPy is in use by falling back to the current aliases in this PR if __skip_array_function__ is not defined.

Let's check with @shoyer to see if this is a reasonable path forward or if there is something else we should be doing.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

__skip_array_function__ is still open numpy/numpy#13585, that means it won't be available at least before NumPy 1.17, which will prevent this PR from getting in until then.

My proposal was to find a solution that would work independent of that. I think __skip_array_function__ is the best solution for this situation, unfortunately not usable yet.

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented May 20, 2019

So I guess we do not need to cover for the case NUMPY_EXPERIMENTAL_ARRAY_FUNCTION='1' and __skip_array_function__ undefined? In that case it seems a reasonable try...except.

I am fine with leaving this open until we have __skip_array_function__ .
To avoid confusion it is probably a good idea to close #2955 though.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

IMO, the test could be skipped for now to get this PR merged. This can be done by checking if NEP-18 is active, just as in https://github.com/dask/dask/blob/master/dask/array/tests/test_array_function.py#L5-L11

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

__skip_array_function__ is still open numpy/numpy#13585, that means it won't be available at least before NumPy 1.17, which will prevent this PR from getting in until then.

That seems like a fair assessment Peter.

My proposal was to find a solution that would work independent of that. I think __skip_array_function__ is the best solution for this situation, unfortunately not usable yet.

Sure. That would amount to resuming PR ( #2955 ) and adding a similar PR for moveaxis. I'm not sure it is worth the effort if we are going to remove it soon anyways. Though wouldn't be opposed to doing it if we decide it is useful.

So I guess we do not need to cover for the case NUMPY_EXPERIMENTAL_ARRAY_FUNCTION='1' and __skip_array_function__ undefined? In that case it seems a reasonable try...except.

IIRC this will become default behavior in NumPy 1.17. So assuming __skip_array_function__ lands in that release. We can just check for that.

I am fine with leaving this open until we have __skip_array_function__ .

Sounds good.

To avoid confusion it is probably a good idea to close #2955 though.

I guess that depends on whether we need a stopgap as Peter alluded to before. Do you have thoughts on this?

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

IMO, the test could be skipped for now to get this PR merged. This can be done by checking if NEP-18 is active, just as in

I'm not sure I completely agree unfortunately. It seems this function will break on different versions of NumPy and we don't have a good way to explain to the user why. Would prefer to see a solution that smoothed over these differences.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

I'm not sure I completely agree unfortunately. It seems this function will break on different versions of NumPy and we don't have a good way to explain to the user why. Would prefer to see a solution that smoothed over these differences.

But having NEP-18 is an opt-in only condition, and the user needs to explicitly set NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1, which is experimental, so they should be aware of issues. However, I don't have a strong opinion, if everybody is fine with keeping this open, that's fine by me too, otherwise if someone is in need of the functionality provided in this PR, we have an alternative to allow things to move forward with a bit of compromise.

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented May 20, 2019

In terms of needing these functions in particular, I suspect the need is fairly limited as in most cases a similar .transpose() function is a possibility. Although I am actually somewhat surprised that these functions are the first where we walk into this NEP-18 issue.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

It's the first time in Dask, but we ran into that in CuPy cupy/cupy#2029, which was actually the trigger to have __skip_array_function__.

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

NEP 18 is still really new. So am not surprised we are running into a few speed bumps. We are all still learning what this means and what it should look like (hence why it is experimental ;). Hopefully we will have this sort of thing worked out by NumPy 1.17 when the protocol most likely will be enabled by default.

@jakirkham

This comment has been minimized.

Copy link
Member

commented May 20, 2019

...if someone is in need of the functionality provided in this PR, we have an alternative to allow things to move forward with a bit of compromise.

I think if someone needs these, we can just say the NumPy functions already work (the gist of this PR). They should try using them as is on Dask Arrays. If that doesn't work, we should try to get a bit more information about why not.

@pentschev

This comment has been minimized.

Copy link
Member

commented May 20, 2019

I think if someone needs these, we can just say the NumPy functions already work (the gist of this PR).

That's a good point. As I mentioned before, I've got no objections in waiting on this PR nor skipping the tests. If we can wait, let's wait.

@jrbourbeau jrbourbeau referenced this pull request Jun 18, 2019
@martindurant

This comment has been minimized.

Copy link
Member

commented Jun 19, 2019

Should this be closed, then, for now? I don't follow quite what the plan is.

@@ -290,6 +290,9 @@ def matmul(a, b):
return out


moveaxis = np.moveaxis

This comment has been minimized.

Copy link
@pentschev

pentschev Jun 20, 2019

Member
Suggested change
moveaxis = np.moveaxis
def numpy_implementation(func, *args, **kwargs):
if hasattr(func, '_implementation'):
return func._implementation(*args, **kwargs)
else:
return func(*args, **kwargs)
moveaxis = partial(numpy_implementation, np.moveaxis)

This comment has been minimized.

Copy link
@pentschev

pentschev Jun 20, 2019

Member

Probably numpy_implementation could go somewhere more general, such as dask/utils.py.

This comment has been minimized.

Copy link
@jakirkham

jakirkham Jun 20, 2019

Member

Seems like a good idea. Could add to dask/array/numpy_compat.py.

This comment has been minimized.

Copy link
@shoyer

shoyer Jun 20, 2019

Member

Please don't use _implementation. This is not a public API exposed by NumPy, as indicated by the name starting with an underscore. (If this isn't clear from the NEP, we should fix that.)

It would be more maintainable to copy the implementations from NumPy. These are not complicated functions, so I really don't think there is much that could go wrong with copying their code.

This comment has been minimized.

Copy link
@pentschev

pentschev Jun 20, 2019

Member

Is there an alternative to copying the implementation? These are not the only cases where this will be necessary. In CuPy as well we need this for some functions that are currently aliases, I've submitted a PR for those cupy/cupy#2249.

Granted, I was very inattentive to _implementation being non-public.

@@ -971,6 +974,9 @@ def roll(array, shift, axis=None):
return result


rollaxis = np.rollaxis

This comment has been minimized.

Copy link
@pentschev

pentschev Jun 20, 2019

Member
Suggested change
rollaxis = np.rollaxis
rollaxis = partial(numpy_implementation, np.rollaxis)
@pentschev

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

I think we shouldn't close this issue, but maybe add [WIP] to the title, just so it becomes visible it can't be merged yet?

NumPy master has now an _implementation attribute for functions, that allows falling back to NumPy even when __array_function__ is enabled. I've added suggestions to replace the aliasing of moveaxis and rollaxis by a dispatcher for NumPy's _implementation.

Once NumPy 1.17 is merged, this change should automatically work there. However, it's important to note that these two functions will never work for NumPy 1.16.x when NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1, but probably when NumPy 1.17 is released, we should use that for the failing test, since it will have better __array_function__ support, which is what we are aiming for with Dask too.

@shoyer

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

@shoyer

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

@pentschev

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

@shoyer I totally missed most of the discussion about that, but I'm assuming _implementation is non-public due to some issue with its use. If it's too complicated to explain briefly, would you mind pointing me to discussion about its consequences and downsides?

@shoyer

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

@jakirkham

This comment has been minimized.

Copy link
Member

commented Jun 20, 2019

Thanks for the feedback, Stephan. In that case, @TAdeJong would you be ok copying the NumPy implementations and putting them both in dask/array/numpy_compat.py?

@jakirkham jakirkham added the array label Jun 20, 2019

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented Jun 21, 2019

I will use some time next week for that.

@pentschev

This comment has been minimized.

Copy link
Member

commented Jun 21, 2019

Thanks for the link @shoyer, that was a good summary. It's really unfortunate that we can't do that for NumPy 1.17, but thanks for all the work you did on that!

@pentschev

This comment has been minimized.

Copy link
Member

commented Jun 21, 2019

@jakirkham @TAdeJong if we copy the implementation, we then will have no dependencies on NumPy 1.17, so we can merge this once that is done.

@TomAugspurger

This comment has been minimized.

Copy link
Member

commented Aug 9, 2019

Where does this PR stand now? We had a user report a warning (I think from NumPy) that move axis wasn't implemented.

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented Aug 9, 2019

I hope to spend some time this weekend to make this mergable.

@TAdeJong

This comment has been minimized.

Copy link
Contributor Author

commented Aug 11, 2019

I now added both implementations from numpy. I added them to array/routines. @jakirkham, was there any specific reason you had thought to put them in numpy_compat?

@jakirkham

This comment has been minimized.

Copy link
Member

commented Aug 11, 2019

Two reasons, we copy other functions there and it contains the proper NumPy license.

@TomAugspurger

This comment has been minimized.

Copy link
Member

commented Aug 12, 2019

Can we also add a test ensuring that np.rollaxis(dask.array.Array) is lazy and returns a Dask Array?

def test_move_axis_numpy():
    a = da.random.random((4, 4, 4), chunks=2)
    result = np.moveaxis(a, 2, 0)
    assert isinstance(result, da.Array)
    assert_eq(result, np.moveaxis(a.compute(), 2, 0))

I can also push that here if you want.

@TomAugspurger
Copy link
Member

left a comment

Pushed to changes

  1. Added the test using the np API
  2. moved the implementation to numpy_compat.

Should be good to go.

@pentschev

This comment has been minimized.

Copy link
Member

commented Aug 12, 2019

LGTM, I would recommend merging it.

Thanks for the work @TAdeJong, and @TomAugspurger for the latest adjustments!

@TomAugspurger TomAugspurger merged commit 244c2fe into dask:master Aug 12, 2019

2 checks passed

continuous-integration/appveyor/pr AppVeyor build succeeded
Details
continuous-integration/travis-ci/pr The Travis CI build passed
Details
@TomAugspurger

This comment has been minimized.

Copy link
Member

commented Aug 12, 2019

Thanks @TAdeJong!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
7 participants
You can’t perform that action at this time.