Simplifying looped np.tensordot expression-CodePudding

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?

CodePudding user response：

You can use numpy.einsum function, in this case

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

CodePudding user response：

An alternative to einsum is matmul/@. The first array has to be transposed so the sum-of-products dimension is last:

In [162]: a = np.random.rand(10,5,2)
     ...: b = np.random.rand(10,5,50)
In [163]: c=a.transpose(0,2,1)@b
In [164]: c.shape
Out[164]: (10, 2, 50)
In [165]: c1 = np.random.rand(10,2,50)
     ...: 
     ...: for i in range(a.shape[0]):
     ...:     c1[i] = np.tensordot(a[i], b[i], axes=(0,0))
     ...: 
In [166]: np.allclose(c,c1)
Out[166]: True

tensordot reshapes and transposes the arguments, reducing the task to simple dot. So while it's fine for switching which axes get the sum-of-products, it doesn't handle batches any better than dot. That's a big part of why matmul was added. np.einsum gives a same power (and more), but performance usually isn't quite as good (unless it's been "optimized" to the equivalent matmul).