python - numpy中3维矩阵的乘法

Question

我想我问了 wrong question昨天。我真正想要的是将两个 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)