numpy中三維矩陣的乘法

Question

我想我昨天問了一個錯誤的問題 。 我真正想要的是2x2xN A兩個2x2xN矩陣A和B

C[:,:,i] = dot(A[:,:,i], B[:,:,i])

例如，如果我有一個矩陣

A = np.arange(12).reshape(2, 2, 3)

如何通過上述定義得到C = A x A ？ 是否有內置函數來執行此操作？

---

此外，如果我將A (shape 2x2xN)與B (shape 2x2x1, instead of N) ，我想得到

C[:,:,i] = dot(A[:,:,i], B[:,:,1])

Answer 1

嘗試使用numpy.einsum ，它有一點學習曲線，但它應該給你你想要的。 這是一個讓你入門的例子。

import numpy as np

A = np.random.random((2, 2, 3))
B = np.random.random((2, 2, 3))

C1 = np.empty((2, 2, 3))
for i in range(3):
    C1[:, :, i] = np.dot(A[:, :, i], B[:, :, i])

C2 = np.einsum('ijn,jkn->ikn', A, B)
np.allclose(C1, C2)