I have two arrays, A and B, with dimensions (l,m,n) and (l,m,n,n), respectively. I would like to obtain an array C of dimensions (l,m,n) which is obtained by treating A and B as matrices in their fourth (A) and third and fourth indices (B). An easy way to do this is:
import numpy as np #Define dimensions l = 1024 m = l n = 6 #Create some random arrays A = np.random.rand(l,m,n) B = np.random.rand(l,m,n,n) C = np.zeros((l,m,n)) #Desired multiplication for i in range(0,l): for j in range(0,m): C[i,j,:] = np.matmul(A[i,j,:],B[i,j,:,:])
It is, however, slow (about 3 seconds on my MacBook). What’d be the fastest, fully vectorial way to do this?
Try to use
It has many use cases, check the docs: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html
Or, for more info, a really good explanation can be also found at: https://ajcr.net/Basic-guide-to-einsum/
In your case, it seems like
should work. Also, you can try the optimize=True flag to get more speed, if needed.
Answered By – Ivelate
Answer Checked By – Candace Johnson (BugsFixing Volunteer)