將 3D 矩陣與 2D 矩陣相乘，返回與 2D 矩陣具有相同維度的矩陣

Question

在 Numpy 中，假設我有一個尺寸為 ixjxk 的矩陣 A，和一個尺寸為 kx i 的矩陣 B。 如何在不使用循環的情況下取回尺寸為 kxi 的產品？
例如：

 Let         A = [[[1 2],
                   [3 4]],
                  [[5 6],
                   [7 8]]]

             B = [[a b],
                  [c d]]

我想得到：

             C = [[a+2c  5b+6d],
                  [3a+4c 7b+8d]]

我目前的解決方案是

np.diagonal(np.dot(A, B), axis1=0, axis2=2)

但是，問題是我正在處理一個大數據集（ A和B有很大的尺寸），所以np.dot(A, B)會導致MemoryError 。 因此，我想找出一種更好的方法來解決這個問題，而無需計算點積。

我已經研究過像einsum和tensordot這樣的函數，但我沒有找到我需要的東西（或者我可能錯過了一些東西）。 如果有人可以幫助我，我將不勝感激。 謝謝！

Answer 1

您可以通過分析輸入和 output 維度來做到這一點。

A.shape -> i, j, k
B.shape -> k, i

你的例子不是很好，但如果你仔細看看你要求 output 是什么， B的形狀必須與A的第一個和最后一個尺寸相匹配。

總和減少發生在A的最后一個軸和B的第一個軸上：

np.einsum('ijk,ki->ji', A, B)

einsum可能令人生畏，但進行此類分析可以為您節省大量因反復試驗而導致的挫敗感。

Answer 2

看着你的 output， C[:,0]是

A[0] @ B[0].T

同樣， C[:,1]是A[1] @ B[1].T 。

考慮到這一點， np.einsum公式為：

C = np.einsum('ijk,ik->ji', A,B)