Turn a Python function into a GUFunc

[mrocklin]

I would like for there to be an easy way to convert a numpy.ndarray -> numpy.ndarray Python function into a generalized universal function (gufunc).

Motivation

In many cases we have a Python function that operates on a multi-dimensional array and produces another multi-dimensional array. As an example, consider the Scikit-Image function, skimage.feature.canny, which consumes and produces a 2d NumPy array.

raw_image = np.random.random((1000, 1000))
processed_image = skimage.feature.canny(raw_image)

If I have a stack of these images I would like for the same function to broadcast across extra dimensions.

raw_stack = np.random.random((10, 1000, 1000))
processed_stack = skimage.feature.canny(raw_stack)  # this doesn't work today

This logic could be placed inside each function, but it would be nice to implement it it once and then decorate functions in the future

@numpy.guvectorize(signature="(n, m) -> (n, m)")  # or something like this
def canny(...):
    ...

GUFuncs

The behavior that I want there is exactly the behavior of GUFuncs, but currently they're hard to construct, except with CPython or with Numba.

In many cases a pure Python decorator that just used for loops would be welcome here. The performance hit of using Python here is relatively low because the function that we're looping over is fairly costly (often 100s of milliseconds). This is the sort of decorator that a project like Scikit-Image (cc @jni and @stefanv ) or ITK (cc @thewtex) would probably be happy to use.

From a Dask perspective this would be great because we can handle gufuncs nicely, rechunking user-provided Dask arrays so that all core dimensions are single-chunked and then automatically and lazily applied across broadcast dimensions.

This conversation came out of the SciPy conference, where @seberg and @mattip discussed it.

[seberg]

We do have another issue (a bit deeper), in that ufuncs currently do not support optional parameter-like kwargs (or optional input arguments, but those are a different story probably), which many functions need.

Aside from that, I think there is a bug with multi-dimensional arrays, and I will dig into it a bit. There is an example in np.vectorize for convolve which seems to work, but the 2-d example seems not to work correctly. (There is also the issue of fixed and optional core dimension not being parsed)

Another thing to keep in mind is that it could be nice to add an external_loop or similar flag which calls the python function with one additional outer dimension for optimization.

EDIT: Oh, the core dimension feature currently is implemented in pure python. We probably should move that into C.

EDIT2: I do not remember what 2-d example I meant, possibly the signature="(n, m) -> (n, m)" from the original post, but that has just too many spaces.

[mrocklin]

This came up again here: dask/dask-blog#27 (comment)

Just for my own planning, is anyone actively working on this? To be clear, I don't expect anything, (you all have plenty to do) I'm just planning out some medium-term work.

[seberg]

Not right now, but since it seems to come up a bit more often. I think we can make sure to pick it up medium-term, there is a reason I tagged it. Maybe when Matti comes back from vacation. (It should not be a huge project really, although touches fairly complicated parts of numpy maybe)

[seberg]

@WarrenWeckesser and @mattip just to ping you on this. Just talked to Matthew again, and ITK is for exmaple insterested in it. So we should probably remember to keep this on the agenda for the intermediate or even near future.

[thewtex]

Friendly ping to help get this on folks radar :-)

[mattip]

note that this is a major blocker for arbitrary functions. Will the ufunc/dtype refactor help solve this?

ufuncs currently do not support optional parameter-like kwargs

[seberg]

Yeah, it dropped from mine, sorry. It may actually be a nice project (for those not scared of spending a larger chunk of time and learning about ufuncs the heads-on way).

ufuncs currently do not support optional parameter-like kwargs

I agree that this is indeed a huge limitation right now in practice, at least if we think of functions like those in skimage.

Will the ufunc/dtype refactor help solve this?

No (Actually yes), but not explicitly: I think it will give us a lot more work with. So in that sense I think it is/was a necessary step before anyone can reasonably approach it. For example, we could push dealing with it into the "inner-loop" UfuncImpl/ArrayMethod. The ArrayMethod.get_loop(...) could be in charge of dealing with such arguments (if it signals that it can).

EDIT: I realized that if we put parameters (read keyword arguments) as options to the inner-loop (ufuncimpl/ArrayMethod), then they cannot modify promotion. That might just be as well, but it means that it may be awkward if you want a parameter change behaviour in certain ways. An example would be np.sqrt(-1., complex=True) is tricky or impossible. Such a thing would always have to work through picking the output domain explicitly: np.sqrt(-1., dtype=np.AbstractComplex). That is probably OK, just something to note :).