einsum equivalent for ndarray multiplication

Question

I have the following multiplication routine:

import numpy as np
a = np.random.rand(3,3)
b = np.random.rand(3,50,50)
res = np.zeros((3, 50, 50))
for i in range(50):
    for j in range(50):
        res[:,i,j] = a @ b[:,i,j]

What is the einsum equivalent expression?

Best regards

Answer 1

Might want to brush up on Einstein summation notation:

res = np.einsum('ij, jkl -> ikl', a, b)

In this case, np.tensordot is also useful:

np.tensordot(a ,b, 1).shape
Out[]: (3, 50, 50)