WIP: ENH: allow __array_ufunc__ to override __matmul__ by mattip · Pull Request #11061 · numpy/numpy

mattip · 2018-05-07T17:55:56Z

The discussion of issue #9028 yielded two different possible directions for allowing ndarray subclasses to override __matmul__ via __array_ufunc__. This is option 2 - wrapping the non-ufunc __matmul__ implementation in a way that respects the __array_ufunc__ conventions.

Here is what changed:

added a ufuncwrapper class that is meant to be used as a decorator, the class itself is implemented in C but the __call__ method is in python for now. Its use is shown in test_wrapper, which shows how it could be used for a method called solve
wrapped array_matrix_multiply with the class, and added the tests at the top of the issue to show how the wrapped method checks __array_ufunc__ before calling matmul.

The wrapper only needs to know nin and nout, which are parameters to the class initializer.

Still todo:

get feedback about this direction, is it what was asked for?
move the python-level wrapper into C (make ufuncwrapper_call work) and clean up code
document ufunc-wrapper
add more tests, make sure there are no refcount leaks
think about edge cases where this will not work, ~~for instance where None is a valid return value from other.__array_ufunc__ (line 397 in number.c)~~ (EDIT: fixed None return via tuple)

Note the other option is to extend the meaning of ufunc signature so it can do multiple dispatch inside get_binary_op_function not only by types but by signature variations as well.

mattip · 2018-05-07T17:57:27Z

commit d96a998 is not part of this pull request, I will remove it in the next iteration

mhvk · 2018-05-08T00:58:31Z

Shouldn't this be *((obj,) + args)? check_override would need to pass on the wrapped function, no?

Sorry, I meant meth (which is the wrapped method, correct?) And then obj will be self of ndarray, which I think would not be passed on at all.

AFAICT check_override does not need meth to call PyUFunc_CheckOverride, it only needs self to provide the nin, nout for argument normalization. The actual test for __array_ufunc__ is in PyUFunc_WithOverride (from src/private/ufunc_override.c) which indeed does need obj, so that line should be

r = self.check_override(*((self, obj) + args), **kwds)

obj is then not different than args, and the code can be simplified to (after taking into account additional corrections)

def __call__(self, meth): def wrap(*args, **kwds): # use status to indicate that one of args overrode the call (status, r) = self.check_override(*((self,) + args), **kwds) if status > 0: return r return meth(*args, **kwds)

I do think meth needs to be passed on, and not the wrapper class as "ufunc", otherwise any class defining __array_ufunc__ has no idea what the real operation is and thus cannot use it to act accordingly. E.g., for astropy's Quantity, I do need to know that the wrapped ufunc is matmul, since that means I know the units of the two arguments have to be multiplied.

But meth is a generic python function or method, not a np.ufunc or np.ufunc_wrapper, specifically it cannot call the bound class method self.check_override.

You get the needed information from func.__name__ in Quantity?

I tried to propagate __name__ and __doc__ in the wrapper __call__, I see I transferred them to the wrapped function wrap but not to self. I will test that those two attributes are available in the ufunc passed in to __array_ufunc__

mhvk · 2018-05-08T01:17:06Z

@mattip - I realize this is still a work in progress but will go ahead and write some first comment anyway:

Mostly to see if I follow the logic: I was slightly wondering about the complication of a new ufunc base when all you seemed to need is nin and nout. At first, I thought perhaps one should just pass those in (this used to happen in the initial implementation for __numpy_ufunc__ -- and this was explicitly for dot and matmul...). But I think this is not enough: in the end (PyObject *)ufunc is passed to __array_ufunc__, so that one at least needs to pass the function as well. And then your approach seems the nicer one (despite the many recasts).
I think None has to be allowed output. One way out might be to put whatever the override produces into a tuple, and then undo the tuple in the python part (I guess this is not all that important, though, since in C this would be trivial to solve).

More general (but I think both can be done after this PR; just have to be sure not to do anything that excludes it):

It is a bit strange to me to have methods rather than functions be wrappable.
One could envisage using your UfuncBase in a new UfuncLike ABC, for which nin and nout must be defined, and __call__, reduce, etc., methods can be overridden by __array_ufunc__.

mattip · 2018-05-08T07:42:16Z

@mhvk Thanks for the feedback

I could theoretically refactor the use of PyUFuncBaseObject to nin, nout, but the signature of __array_ufunc__(self, ufunc, method, *args, **kwargs) would have to change, which would be too drastic a change
tuple output to check_override - nice, adopted.
the comments in PyUFunc_CheckOverride expressly mention "unbound method" but indeed, it seems the code would work with a function.

mattip · 2018-05-08T12:26:54Z

It turns out using __array_ufunc__ on an old-style class

class OldClass():
    def __array_ufunc__(...):
        ....

on python2 is not supported. I added a warning

mhvk · 2018-05-08T16:49:16Z

One more major point (though perhaps I misunderstand): so far, __array_ufunc__ implementations can count on being passed an actual ufunc as the first argument, and at least in my use for astropy's Quantity, what gets done depends on it (specifically, a helper function is gotten from a dict keyed by ufuncs, which calculates the unit of the result from the input units, and does possible unit conversion on the input). I think in general, implementations need to be able to determine what the function is that is being overridden, i.e., __array_ufunc__ needs to be passed (some form of) np.matmul (so that, e.g., Quantity knows to multiply the units of the two input arguments together).

As far as I understand, though, in your current implementation, UfuncWrapper returns a wrapped solve which when executed will end up in ufunc_check_override with self being a UfuncWrapper instance. The latter is then recast as a PyUFuncBaseObject and passed on to PyUFunc_CheckOverride. Hence, any __array_ufunc__ will not get solve, but rather a UfuncWrapper instance.

mhvk · 2018-05-08T17:06:49Z

Sorry, I missed your earlier reply in the in-code comments before sending the above comment. In principle, one could indeed use __name__, but I think the main thing to guarantee is uniqueness of the ufunc argument in __array_ufunc__. This is partially why I was wondering if np.matmul itself could become a UfuncWrapper instance...

mhvk · 2018-05-08T18:59:33Z

@mattip - just so you have a reference to an actual __array_ufunc__ implementation:

It would be very helpful to ensure that functions passed on to __array_ufunc__ are identical to actual functions being exposed in the numpy namespace.

mattip · 2018-05-08T20:21:10Z

I see. Need to think about this a bit.

I am working on another WIP PR where matmul becomes a full-fledged ufunc without a wrapper, at the price of a more complicated signature and needing to enhance some other functions in the ufunc machinery. If I can do it, we can compare and decide which seems more palatable.

mattip · 2018-05-29T05:13:33Z

It seems this approach will not be pursued now, in favor of making matmul a true ufunc via expanding the core signature. Closing.

mhvk reviewed May 8, 2018

View reviewed changes

mattip force-pushed the ufunc_wrapper branch 4 times, most recently from 645f175 to b6f69bb Compare May 8, 2018 12:25

mattip force-pushed the ufunc_wrapper branch from b6f69bb to 47fdc9f Compare May 8, 2018 13:03

charris changed the title ~~WIP: BUG: allow __array_ufunc__ to override __matmul__~~ WIP: ENH: allow __array_ufunc__ to override __matmul__ May 9, 2018

charris added 01 - Enhancement component: numpy._core labels May 9, 2018

This was referenced May 22, 2018

WIP: ENH: Expand gufunc signature to allow flexible dimension specs #11132

Closed

WIP: ENH: Matmul as ufunc #11133

Closed

mattip added 6 commits May 23, 2018 10:18

define PyUFuncBaseObject and use in PyUFunc_CheckOverride

7a9877a

add ufunc_wrapper object that decorates ufunc-like class methods

d9c07aa

add c- and python- level check_override

936153b

check_override is a bound class method, no self needed. Returns tuple

40b5b79

test, add warning for old-style classes in python2

8154bcd

test, fix for __name__ and __doc__ on ufunc

b2f17fd

mattip force-pushed the ufunc_wrapper branch from ec9b824 to b2f17fd Compare May 23, 2018 17:46

mattip closed this May 29, 2018

mattip mentioned this pull request Jul 8, 2018

Implement gufunc(pyfunc, ...) similar to numpy.vectorize but without the vectorize part #10526

Open

mattip deleted the ufunc_wrapper branch November 15, 2018 18:29

Conversation

mattip commented May 7, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattip commented May 7, 2018

Uh oh!

mhvk May 8, 2018

Choose a reason for hiding this comment

Uh oh!

mhvk May 8, 2018

Choose a reason for hiding this comment

Uh oh!

mattip May 8, 2018

Choose a reason for hiding this comment

Uh oh!

mhvk May 8, 2018

Choose a reason for hiding this comment

Uh oh!

mattip May 8, 2018

Choose a reason for hiding this comment

Uh oh!

mhvk commented May 8, 2018

Uh oh!

mattip commented May 8, 2018

Uh oh!

mattip commented May 8, 2018

Uh oh!

mhvk commented May 8, 2018

Uh oh!

mhvk commented May 8, 2018

Uh oh!

mhvk commented May 8, 2018

Uh oh!

mattip commented May 8, 2018

Uh oh!

mattip commented May 29, 2018

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mattip commented May 7, 2018 •

edited

Loading