如何有效地计算numpy中两个矩阵级数的外积？

Question

假设我有A（KxMxN）和B（KxLxN）矩阵，其中L，M，N都很小，K是一个大数。我想计算使用最后2个维度沿着第一个维度得到矩阵C（KxMxL）的外积。在

我可以为“k”中的每个索引运行for循环，并对2D矩阵使用numpy的matmul函数

out = [np.matmul(x,y.T) for x, y in zip(A, B)]
out=np.asarray(out)

我不知道是否可以不使用for循环/理解，因为K是一个非常大的数字。在

Answer 1

matmul起作用，在B中有一个转置，因此其倒数第二个dim与A中的最后一个相匹配。在

In [1019]: A=np.random.rand(K,M,N)
In [1021]: B=np.random.rand(K,L,N)

In [1023]: C=np.einsum('kmn,kln->kml',A,B)
In [1024]: C.shape
Out[1024]: (2, 4, 3)

In [1026]: D=A@B.transpose(0,2,1)
In [1027]: D.shape
Out[1027]: (2, 4, 3)

In [1028]: np.allclose(C,D)
Out[1028]: True

对于这个小例子，timeit是相同的。在

[np.dot(x,y.T) for x, y in zip(A, B)]执行相同的操作；将y的最后第二个维度与{}的最后一个维度匹配，并迭代{}和{}的第一个维度。在

Answer 2

由于A具有形状(K, M, N)，而{}具有形状{}，并且您希望找到形状为(K, M, L)的积和，可以使用np.einsum：

C = np.einsum('kmn,kln->kml', A, B)