Search code examples
pythonnumpymatrixvectorization

Sum of many outer products in NumPy

I have two matrices A and B both with shape (N,M). I would like to perform the following operation: C = np.sum(A[:,None,:]*B[:,:,None],axis=(1,2)), which corresponds to performing the sum of the outer product of each row of A with each row of B. C would then have a shape of (N,).

The problem is that I get a MemoryError when using this form because N=12000 and M=4000.

Is there a way to perform this operation without having to first build the (huge) intermediary array to be summed ?

I suspect a solution with np.einsum would do the trick, but I'm not familiar with it !

Solution

  • Not sure if np.einsum solves the memory issue, but below is equivalent of your calculation using it:

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