I have an application where I have to do a lot of (costly) matrix multiplication. Many of these matrices, but not all, have special structure for which it is easy for me to define a custom data representation and a custom matrix multiplication routine. I want to be able to mix and match these with arbitrary matrices that are represented as ordinary 2d ndarrays.
To take the simplest possible example, I would like to do something like this:
import numpy as np
class IdentityMatrix:
def __matmul__(self, other):
return other
def __rmatmul__(self, other):
return other
__array_priority__ = 10000
where an IdentityMatrix knows that it does nothing when acting upon a matrix on either the left or the right. My hope was that setting __array_priority__ to a high number would cause this to always override an ndarray's matrix multiplication when the other argument is an ndarray, but this does not work:
In [2]: A = IdentityMatrix
...: B = np.array([[3,4],[6,1]])
...: A @ B
...:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-2-8f0beccd3079> in <module>()
1 A = IdentityMatrix
2 B = np.array([[3,4],[6,1]])
----> 3 A @ B
ValueError: matmul: Input operand 0 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1)
Reversing the inputs gives the same error. The frustrating thing is that this does seem to work for operations other than matmul.
Is there any simple way to override matmul in the way I'd like?
Thanks!
