All geek questions in one place

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

Suppose I have matrices A (KxMxN) and B (KxLxN), where L, M, N are small, and K is a large number. I would like to calculate the external product using the last two dimensions along the first dimension to get the matrix C (KxMxL).

I can do this by running a for loop for each index k in "K" and using the numpy matmul function for two-dimensional matrices

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

I wonder if I can do this without a loop / understanding, since K is a very large number.

+5
source share
2 answers

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

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

matmul works with transposition in B , so its second to last dim corresponds to 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]: D=A@B.transpose (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, timeit same.

[np.dot(x,yT) for x, y in zip(A, B)] does the same; match the second last dull y with the symbol x and iterate over the 1st dull A and B

0
source

Source: https://habr.com/ru/post/1263408/

More articles:


All Articles