在numpy中进行三维矩阵相乘

Question

我觉得我昨天问的问题有点偏了。我其实想要的是把两个 2x2xN 的矩阵 A 和 B 相乘，这样做。

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

举个例子，如果我有一个矩阵

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

我该怎么得到 C = A x A，按照上面的定义？有没有现成的函数可以做到这一点？

---

另外，如果我把 A (形状是 2x2xN) 和 B (形状是 2x2x1，而不是 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)