python - NumPy 中许多外积的总和

Question

我有两个矩阵 A 和 B，它们的形状都是 (N,M)。我想执行以下操作:C = np.sum(A[:,None,:]*B[:,:,None],axis=(1,2))，其中对应于A的每一行与B的每一行的外积求和。 C 将具有 (N,) 的形状。

问题是我在使用此表单时遇到了 MemoryError，因为 N=12000 和 M=4000。

有没有一种方法可以执行此操作而不必首先构建要求和的(巨大的)中间数组？

我怀疑使用 np.einsum 的解决方案可以解决问题，但我不熟悉它!

Answer 1

不确定 np.einsum 是否解决了内存问题，但下面等效于您使用它的计算:

C = np.einsum('ij,ik->i',A,B)