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?

CodePudding user response：

Try to use einsum.

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

np.einsum('dhi,dhij->dhj',A,B)

should work. Also, you can try the optimize=True flag to get more speed, if needed.