Почему матрица умножения NumPy работает в одном направлении, а не в транспонированном направлении?

Question

Рассмотрим следующий матричный продукт между двумя массивами:

import numpy as np
A = np.random.rand(2,10,10)                                             
B = np.random.rand(2,2)                                                 
C = A.T @ B

... идет нормально. Я думаю о вышесказанном как о векторном матричном произведении 1 на 2, умноженном на 2 на 2, транслируемом по 2-му и 3-му измерениям 10 на 10 А. Проверка результата C подтверждает эту интуицию;np.allclose(C[i,j], A.T[i,j] @ B) для всех i, j.

Теперь математически я смогу вычислить C.T, а также: B.T @ A, но:

B.T @ A                                                                
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-32-ffdbb14ca160> in <module>
----> 1 B.T @ A

ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 10 is different from 2)

Таким образом, с точки зрения широковещания тензор 10 на 10 на 2 и матрица 2 на 2 совместимы по отношению к матричному произведению, но матрица 2 на 2 и 2 на 10 на 10тензор не так ли?

Информация о бонусе: Я хочу иметь возможность вычислять "квадратичный продукт" A.T @ B @ A, и меня действительно раздражает необходимость писать циклы for для передачи вручную«по одному из измерений. Такое ощущение, что должно быть возможно сделать это более элегантно. У меня большой опыт работы с Python и NumPy, но я редко перехожу за пределы двумерных массивов.

Что мне здесь не хватает? Есть ли что-то в том, как transpose работает с тензорами в NumPy, что я не понимаю?

Answer 1

In [194]: A = np.random.rand(2,10,10)                                           
     ...:    
     ...: B = np.random.rand(2,2)                                               
In [196]: A.T.shape                                                             
Out[196]: (10, 10, 2)

In [197]: C = A.T @ B                                                           
In [198]: C.shape                                                               
Out[198]: (10, 10, 2)

Эквивалент einsum:

In [199]: np.allclose(np.einsum('ijk,kl->ijl',A.T,B),C)                         
Out[199]: True

или включение транспонирования в индексирование:

In [200]: np.allclose(np.einsum('kji,kl->ijl',A,B),C)                           
Out[200]: True

Обратите внимание, что k является суммированным измерением. j и l - другие dot размеры. i является своего рода «пакетным» измерением.

Или, как вы объясните np.einsum('k,kl->l', A.T[i,j], B)

Чтобы получить C.T, индексы результата einsum должны быть lji, илиlk,jki->lji:

In [201]: np.allclose(np.einsum('lk,jki->lji', B.T, A.transpose(1,0,2)), C.T)      
Out[201]: True

In [226]: np.allclose(np.einsum('ij,jkl->ikl', B.T, A), C.T)                       
Out[226]: True

Для сопоставления [201] с @ требуется дополнительная транспонирование:

In [225]: np.allclose((B.T@(A.transpose(1,0,2))).transpose(1,0,2), C.T)          
Out[225]: True

С einsum, когда можно расположить оси в любом порядке, но сmatmul, порядок фиксирован (batch, i, k)@(batch, k, l) -> (batch, i, l) (где могут передаваться измерения batch).

Ваш пример может быть проще, если A имел форму (2,10,9) и B (2,3), с C, в результате (9,10,3)

In [229]: A = np.random.rand(2,10,9); B = np.random.rand(2,3)                   
In [230]: C = A.T @ B                                                           
In [231]: C.shape                                                               
Out[231]: (9, 10, 3)
In [232]: C.T.shape                                                             
Out[232]: (3, 10, 9)

In [234]: ((B.T) @ (A.transpose(1,0,2))).shape                                    
Out[234]: (10, 3, 9)
In [235]: ((B.T) @ (A.transpose(1,0,2))).transpose(1,0,2).shape                   
Out[235]: (3, 10, 9)
In [236]: np.allclose(((B.T) @ (A.transpose(1,0,2))).transpose(1,0,2), C.T)        
Out[236]: True