Переопределить matmul для специальных матриц

Question

У меня есть приложение, в котором мне приходится много (дорогостоящего) умножения матриц. Многие из этих матриц, но не все, имеют специальную структуру, для которой мне легко определить пользовательское представление данных и пользовательскую процедуру умножения матриц. Я хочу иметь возможность смешивать и сопоставлять их с произвольными матрицами, которые представлены как обычные 2d ndarrays.

Чтобы взять самый простой пример, я хотел бы сделать что-то вроде этого:

import numpy as np

class IdentityMatrix:
    def __matmul__(self, other):
        return other

    def __rmatmul__(self, other):
        return other

    __array_priority__ = 10000

, где IdentityMatrix знает, что он ничего не делает, воздействуя на матрицу слева или справа. Я надеялся, что установка __array_priority__ на большое число приведет к тому, что это всегда переопределит умножение матрицы ndarray, когда другой аргумент является ndarray, но это не работает:

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)

Реверсирование входов дает ту же ошибку. Огорчает то, что , кажется, работает для операций, отличных от matmul .

Есть ли какой-нибудь простой способ переопределить матмула так, как я хотел бы?

Спасибо!