CODEWITHSUNDEEP
pythonnumpymatrix
genezin
genezin

Reputation: 63

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

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.

Upvotes: 5

Views: 509

Answers (2)

hpaulj
hpaulj

Reputation: 231550

matmul works, with a transpose in B, so its 2nd to last dim matches the last of A.

In [1019]: A=np.random.rand(K,M,N)
In [1021]: B=np.random.rand(K,L,N)

In [1023]: C=np.einsum('kmn,kln->kml',A,B)
In [1024]: C.shape
Out[1024]: (2, 4, 3)

In [1026]: [email protected](0,2,1)
In [1027]: D.shape
Out[1027]: (2, 4, 3)

In [1028]: np.allclose(C,D)
Out[1028]: True

For this small example the timeit is the same.

[np.dot(x,y.T) for x, y in zip(A, B)] does the same thing; match the 2nd last dim of y with that of x, and iterate on the 1st dim of A and B.

Upvotes: 0

unutbu
unutbu

Reputation: 880469

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)

Upvotes: 5

Related Questions