Vectorized computation of numpys tensor dot

Question

I have two vectors containing tensors of shape (3,3) and shape (3,3,3,3) respectively. The vectors have the same length, I am computing the element-wise tensor dot of these two vectors . For example, want to vectorise the following computation to improve performance:

a = np.arange(9.).reshape(3,3)
b = np.arange(81.).reshape(3,3,3,3)
c = np.tensordot(a,b)

a_vec = np.asanyarray([a,a])
b_vec = np.asanyarray([b,b])    
c_vec = np.empty(a_vec.shape)

for i in range(c_vec.shape[0]):
    c_vec[i, :, :] = np.tensordot(a_vec[i,:,:], b_vec[i,:,:,:,:])

print(np.allclose(c_vec[0], c))
# True

I thought about using numpy.einsum but can't figure out the correct subscripts. I have tried a lot of different approaches but failed so far on all of them:

# I am trying something like this
c_vec = np.einsum("ijk, ilmno -> ijo", a_vec, b_vec)

print(np.allclose(c_vec[0], c))
# False

But this does not reproduce the iterative computation I want above. If this can't be done using einsum or there is a more performant way to do this, I am open for any kind of solutions.

Answer 1

tensor_dot has an axes argument you can use too:

c_vec = np.tensordot(a_vec, b_vec, axes=([1, 2], [1, 2]))

Answer 2

Vectorized way with np.einsum would be -

c_vec = np.einsum('ijk,ijklm->ilm',a_vec,b_vec)