[Numpy-discussion] Allowing broadcasting of code dimensions in generalized ufuncs

Eric Wieser wieser.eric+numpy at gmail.com
Tue Jun 12 02:59:36 EDT 2018

I don’t understand your alternative here. If we overload np.matmul using
*array_function*, then it would not use *ether* of these options for
writing the operation in terms of other gufuncs. It would simply look for
an *array_function* attribute, and call that method instead.

Let me explain that suggestion a little more clearly.

   1. There’d be a linalg.matmul2d that performs the real matrix case,
   which would be easy to make as a ufunc right now.
   2. __matmul__ and __rmatmul__ would just call np.matmul, as they
   currently do (for consistency between np.matmul and operator.matmul,
   needed in python pre- at -operator)
   3. np.matmul would be implemented as:

   @do_array_function_overridesdef matmul(a, b):
       if a.ndim != 1 and b.ndim != 1:
           return matmul2d(a, b)
       elif a.ndim != 1:
           return matmul2d(a, b[:,None])[...,0]
       elif b.ndim != 1:
           return matmul2d(a[None,:], b)
           # this one probably deserves its own ufunf
           return matmul2d(a[None,:], b[:,None])[0,0]

   4. Quantity can just override __array_ufunc__ as with any other ufunc
   5. DataArray, knowing the above doesn’t work, would implement something

   @matmul.register_array_function(DataArray)def __array_function__(a, b):
       if a.ndim != 1 and b.ndim != 1:
           return matmul2d(a, b)
           # either:
           # - add/remove dummy dimensions in a dataarray-specific way
           # - downcast to ndarray and do the dimension juggling there

Advantages of this approach:


   Neither the ufunc machinery, nor __array_ufunc__, nor the inner loop,
   need to know about optional dimensions.

   We get a matmul2d ufunc, that all subclasses support out of the box if
   they support matmul

-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20180611/9e8d6f70/attachment.html>

More information about the NumPy-Discussion mailing list