How to efficiently calculate the outer product of two series of matrices in numpy?

Question

Lets say I have A(KxMxN) and B(KxLxN) matrices, where L,M,N are small and K is a large number. I would like to calculate the outer product of using the last 2 dimensions along the 1st dimension to get matrix C(KxMxL).

I can do this by running a for loop for each index k in "K" and use numpy's matmul function for 2D matrices

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

I wonder if I can do it without the for loop/comprehension as K is a very large number.

Answer 1

Since A has shape (K, M, N) and B has shape (K, L, N), and you wish to find the sum of products with shape (K, M, L), you could use np.einsum:

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