tensordot equivalent of numpy-einsum

Question

I am trying to figure out the tensordot equivalent of the following expression, as sparse package does not support einsum (the sparseness of the original problem is much better than the example below).

The place I am stuck with is the contraction ii->i, I don't know how to interpret this.

mki_shape=(25,25,121)
mki=np.random.uniform(size=mki_shape)
tik_shape=(10,121,25)
tik=np.random.uniform(size=tik_shape)

tim=np.einsum('mki,tik->tim',mki,tik)
print(tim.shape)

Answer 1

You can achieve the result as with np.einsum with a for loop, np.matmul and a np.transpose operation:

mki_shape=(25,25,121)
mki=np.random.uniform(size=mki_shape)
tik_shape=(10,121,25)
tik=np.random.uniform(size=tik_shape)

tim=np.einsum('mki,tik->tim',mki,tik)
print(tim.shape)
(10, 121, 25)

tim2 = np.array([np.matmul(mki[:,:,i],tik[:,i,:].T) for i in range(mki_shape[2])])
tim2 = np.transpose(tim2,axes=(2,0,1))
print(tim2.shape)
(10, 121, 25)

np.allclose(tim,tim2)
True

All you need to know is the dimension which you want to reduce and take care that the dimensions for matmul line up (hence the transpose).