Search code examples
pythonnumpytensordot

Simplifying looped np.tensordot expression

Currently, my script looks as follows:

import numpy as np

a = np.random.rand(10,5,2)
b = np.random.rand(10,5,50)
c = np.random.rand(10,2,50)

for i in range(a.shape[0]):
    c[i] = np.tensordot(a[i], b[i], axes=(0,0))

I want to replicate the same behaviour without using a for loop, since it can be done in parallel. However, I have not found a neat way yet to do this with the tensordot function. Is there any way to create a one-liner for this operation?

Solution

  • You can use numpy.einsum function, in this case

    c = np.einsum('ijk,ijl->ikl', a, b)