How to vectorize multiple matrix multiplication-CodePudding

I have a 2d matrix A[1000*90] and B[90*90*1000]

I would like to calculate C[1000*90]


For i in range(1000)
  C[i,:]=np.matmul(A[i,:],B[:,:,i]

I understand if I use a vectorized formula it's going to be faster, seems like einsum might be the function I'm looking for, but I am having trouble cyphering the syntax of einsum. Is it np.einsum(ij,jki->ik,A,B)?

CodePudding user response：

Your einsum is correct. But there is a better way as pointed out by hpaulj.

Using Matmul:

import numpy as np
A =np.random.rand(1000,90)
B = np.random.rand(90,90,1000)
C = A[:,np.newaxis,:]@B.transpose(2,0,1) ## Matrix multiplication
C = C = C.reshape(-1,C.shape[2])
np.array_equal(C,np.einsum('ij,jki->ik',A,B)) # check if both give same result