pythonpython-3.xnumpynumpy-ndarraynumpy-einsum

einsum equivalent for ndarray multiplication

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

Solution

  • 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)