Sum of many outer products in NumPy

user_7565410 0
Viewed 97

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 !

vectorization matrix numpy python
2 Answers

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)
user_49759810

For future readers: this operation is mathematically equivalent to multiplying the sum of each row of A with the sum of each row of B. In other words, the fastest solution is C = np.sum(A,axis=-1) * np.sum(B,axis=-1).

This is much faster than np.einsum and any other way of computing it.

user_75654100
Related