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

Question

假设我有A(KxMxN)和B(KxLxN)矩阵,其中L,M,N很小而K是一个大数.我想计算沿第一维使用最后2维的外积来得到矩阵C(KxMxL).

我可以通过为"K"中的每个索引k运行for循环并使用numpy的matmul函数用于2D矩阵来完成此操作

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

我想知道如果没有for循环/理解我是否可以做到,因为K是一个非常大的数字.

Answer 1

由于A具有形状(K, M, N)和B形状(K, L, N),并且您希望找到具有形状的产品的总和(K, M, L),您可以使用np.einsum:

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