Computação vetorizada do ponto tensor numpys

Question

Eu tenho dois vetores contendo tensores de forma (3,3)e forma, (3,3,3,3)respectivamente. Os vetores têm o mesmo comprimento, estou calculando o ponto tensor elemento-a-elementar desses dois vetores. Por exemplo, deseja vetorizar a seguinte computação para melhorar o desempenho:

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

Eu pensei em usar numpy.einsum, mas não consigo descobrir os subscritos corretos. Eu tentei muitas abordagens diferentes, mas falhei até agora em todas elas:

# 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

Mas isso não reproduz a computação iterativa que desejo acima. Se isso não puder ser feito usando o einsum ou houver uma maneira mais eficiente de fazer isso, estou aberto a qualquer tipo de solução.

Answer 1

Maneira vetorizada com np.einsumseria -

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

Answer 2

tensor_dottem um axesargumento que você pode usar também:

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