BUG: no way to override matmul/@ if __array_ufunc__ is set

[shoyer]

With NumPy master, it is currently impossible to override ndarray @ other if other also implements __array_ufunc__.

Consider:

import numpy as np

np.__version__  # '1.13.0.dev0+9ef5891'

class OptOut:
    __array_ufunc__ = None
    def __rmatmul__(self, other):
       return 'rmatmul'

class OtherArray:
    def __array_ufunc__(self, *args, **kwargs):
        return 'array_ufunc'
    def __rmatmul__(self, other):
        return 'rmatmul'

array = np.arange(3)
opt_out = OptOut()
other_array = OtherArray()

# OptOut works as expected:
array * opt_out  # TypeError: unsupported operand type(s) for *: 'numpy.ndarray' and 'OptOut'
array @ opt_out  # 'rmatmul'

# But OtherArray does not:
array * other_array  # 'array_ufunc'
array @ other_array  # TypeError: Object arrays are not currently supported

This is problematic. We definitely need a way to allow for overriding @ from the second argument, whether than means breaking the rules on __array_ufunc__ (allowing matmul to be passed even though it isn't a ufunc) or breaking the rules for __matmul__ (to call __rmatmul__ or allow __rmatmul__ to be called by Python, even when __array_ufunc__ is implemented on the other argument).

[shoyer]

Probably the simplest thing to do is to make np.matmul use __array_ufunc__ again, even though it isn't a real ufunc. It's probably less surprising than for @ not be overridable if you use __array_ufunc__.

[njsmith]

The whole reason binops care about __array_ufunc__ is that they dispatch to ufuncs, so saying that binops that don't dispatch to ufuncs don't pay attention to __array_ufunc__ seems pretty defensible. OTOH I worry that if we do that then it'll become some big breaking change if/when matmul starts dispatching to a ufunc...

[shoyer]

The whole reason binops care about array_ufunc is that they dispatch to ufuncs, so saying that binops that don't dispatch to ufuncs don't pay attention to array_ufunc seems pretty defensible.

Agreed, this is also totally reasonable. __matmul__ could continue the current state of affairs, deferring based on __array_priority__.

OTOH I worry that if we do that then it'll become some big breaking change if/when matmul starts dispatching to a ufunc...

Yes, this would be my concern as well. We only get one chance to introduce __array_ufunc__ for the first time.

[charris]

The @ operator currently forwards on __array_priority__ and __array_ufunc__ = None.

[charris]

I think the opt-out is dual use, it basically tells arrays to "leave me alone". If arrays weren't such greedy bastards and only ate other arrays, we wouldn't need this solution.

[mhvk]

Arguably easiest would be to get matmul become a gufunc -- it didn't seem terribly difficult, but might be less efficient.

[eric-wieser]

Alternatively, could we just create some ufunc_like class, that exposes nout and nin? That should cover enough cases that we can just pass np.matmul into __array_ufunc__, even if it's not really a ufunc.

[mhvk]

Note that the basic matrix multiply gufunc is already written: https://github.com/numpy/numpy/blob/master/numpy/core/src/umath/umath_tests.c.src#L143

What would seem to be needed is:

• pre/append appropriate 1s in the dimension for vector-matrix and matrix-vector, or, probably better, call specific appropriate gufuncs for those since I guess can cannot guarantee that all operands can do appropriate reshaping (I vaguely recall we had a discussion about names...).
• special-case the blas call inside the gufunc loop for matrix times matrix

@seberg, @eric-wieser - since you last looked at that file -- is this possible?

[shoyer]

I think it would be great to have matmul be a gufunc, if/when gufuncs support multiple dispatch. But I'm skeptical about dispatching to multiple gufuncs for np.matmul.

If you write ndarray @ other, how can ndarray even know which of matmat_multiply, matvec_multiply or vecmat_multiply to pass to other.__numpy_ufunc__, without inspecting potentially opaque details about other? For example, if other is an xarray.Dataset (i.e., a dict of labeled arrays), other.ndim is not even defined.

[charris]

I don't think the call needs to go directly to the gufunc, there can be preparation before and after the call, the override should still work as long as there is only one call to the gufunc. And the return can be reshaped. Multiple calls would get dicey.

[charris]

BTW, I don't think this needs to be in 1.13, we already have too much stuff.

[mhvk]

But it would be good to solve the issue here one way or another, of not being able to override __matmul__. My own sense would still be that separate matmat, matvec and vecmat gufuncs might still be the easiest route; they'd be useful even apart from helping solve matmul.

[charris]

Plus vecvec. In general, one would like to reshape (..., i, j) dot (j,) to (n,j) dot (j,) for efficiency in multiplication, and then unreshape, same with various other combinations, plus BLAS has contiguity requirements, so some preliminary work is always required. Timewise, there doesn't seem to be much difference between the special cases and reshaping and having just one function may make things easier for the user who wants to override.