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

Question

假设我有 A(KxMxN) 和 B(KxLxN) 矩阵，其中 L、M、N 很小，K 很大。我想计算使用沿第一个维度的最后两个维度得到矩阵 C(KxMxL) 的外积。

我可以通过为“K”中的每个索引 k 运行 for 循环并使用 numpy 的二维矩阵 matmul 函数来实现此目的

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)