WIP: ENH: allow __array_ufunc__ to override __matmul__ #11061
Conversation
|
commit d96a998 is not part of this pull request, I will remove it in the next iteration |
numpy/core/numeric.py
Outdated
| ''' | ||
| def __call__(self, meth): | ||
| def wrap(obj, *args, **kwds): | ||
| r = self.check_override(*((self,) + args), **kwds) |
There was a problem hiding this comment.
Shouldn't this be *((obj,) + args)? check_override would need to pass on the wrapped function, no?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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__
|
@mattip - I realize this is still a work in progress but will go ahead and write some first comment anyway:
More general (but I think both can be done after this PR; just have to be sure not to do anything that excludes it):
|
|
@mhvk Thanks for the feedback
|
645f175 to
b6f69bb
Compare
|
It turns out using on python2 is not supported. I added a warning |
|
One more major point (though perhaps I misunderstand): so far, As far as I understand, though, in your current implementation, |
|
Sorry, I missed your earlier reply in the in-code comments before sending the above comment. In principle, one could indeed use |
|
@mattip - just so you have a reference to an actual
It would be very helpful to ensure that functions passed on to |
|
I see. Need to think about this a bit. I am working on another WIP PR where |
|
It seems this approach will not be pursued now, in favor of making matmul a true ufunc via expanding the core signature. Closing. |
The discussion of issue #9028 yielded two different possible directions for allowing
ndarraysubclasses 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:
ufuncwrapperclass 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 intest_wrapper, which shows how it could be used for a method calledsolvearray_matrix_multiplywith the class, and added the tests at the top of the issue to show how the wrapped method checks__array_ufunc__before callingmatmul.The wrapper only needs to know
ninandnout, which are parameters to the class initializer.Still todo:
ufuncwrapper_callwork) and clean up codeufunc-wrapperfor instance where(EDIT: fixedNoneis a valid return value fromother.__array_ufunc__(line 397 in number.c)Nonereturn via tuple)Note the other option is to extend the meaning of ufunc signature so it can do multiple dispatch inside
get_binary_op_functionnot only by types but by signature variations as well.