大量矩阵对的有效乘法

Question

我有两组10,000个矩阵存储为两个ndarrays A和B：

A.shape = (10000, 2, 10)
B.shape = (10000, 10, 2)

有多种方法可以将这些矩阵相乘以获得：

C.shape = (10000, 2, 2)

与C[k,:,:] = dot(A[k,:,:], B[k,:,:])？

Answer 1

您可以使用numpy.einsum：

np.einsum('ijk,ikl->ijl', A, B)