Support for Generalized Universal Functions

[emcastillo]

This is a big nasty one.

Currently wip.

Adds a GUFunc class (all the axes and batching stuff has been copied straight from dask).

There are some heavy incompatibilities here with things such as matmul since all the batching is done in the relevant cublas routine, we should avoid doing repeated calls at any cost. So for some gufuncs we need to skip several of these steps.

Also, we need to find a way to blend this with the regular UniversalFunction, probably wrapping it and doing a straight call should be enough, but we will provide additional support for currently non-supported kwargs such as where.

This needs to be cythonized too.

Fixes:
#4549
#2552

[jakirkham]

cc @seberg @pentschev (in case this is of interest 🙂)

[leofang]

Since @seberg its pinged, I'd like to know whether #4549 (which this PR is trying to fix) should be considered as a NumPy bug or not. NumPy made the gufuncs not dispatchable via __array_function__ IMHO, see #4549 (comment).

[seberg]

I am not sure I understand. Since __array_ufunc__ takes complete control, it doesn't matter whether the cupy version of matmul is implemented as a ufunc? IIRC the __array_function__ NEP basically says, __array_ufunc__ is preferable, and any function may be moved to that (i.e. be converted into a ufunc if applicable).

I admit it is a bit weird that some cases will actually dispatch to a ufunc from within a non-ufunc wrapped with __array_function__

[emcastillo]

Hi!
in CuPy right now matmul is a regular function instead of a ufunc and does not support most of the ufunc kwargs, so we are not dispatching properly. The issue is entirely in CuPy function dispatcher.
This PR is adding a wrapper to be used by matmul and other ufuncs so that now dispatch will behave properly, and we also get arguments such as axes and others in routines that did not support them before.

[emcastillo]

Btw, although the code is 60% straight from dask I added support for stuff such as optional dimensions in signatures (m?, n), (n, k?) -> (m?, k?), broadcasting, out kwarg, and now I am working on supporting dtypes and type signatures. So hopefully we will have a numpy 1:1 GUFunc API.

How should I reference Dask here? All the matrix transposition credits goes to them, so I don't know if I have to reproduce their license in the code.

[pentschev]

How should I reference Dask here? All the matrix transposition credits goes to them, so I don't know if I have to reproduce their license in the code.

Usually Dask itself adds a comment to where it was taken from with a GitHub permalink (see https://github.com/dask/dask/blob/91cee5b7944e47959f475eaebc42dc75be14bc21/dask/array/numpy_compat.py#L167-L208 for reference), so maybe do the same here?

[emcastillo]

PTAL,
Still some work left for future PRs.

Cythonize -> better performance
Improve test coverage
Support multiple outputs
Support dtype, order and the rest of kwargs.

BTW, matmul dispatch works now, so Dask will not complain :)

[leofang]

I haven't finished a full pass and will do so asap 😅 Found two minor issues so far.

[emcastillo]

PTAL!

[asi1024]

@emcastillo Could you merge the master branch to pass CI?

[emcastillo]

rebased :)

[asi1024]

Jenkins, test this please.

[chainer-ci]

Jenkins CI test (for commit b1bce6a, target branch master) failed with status FAILURE.

[leofang]

Jenkins, test this please

[chainer-ci]

Jenkins CI test (for commit b1bce6a, target branch master) succeeded!

[asi1024]

LGTM!