Simplify pytorch einsum

Question

Consider the following pytorch snippet:

X = torch.einsum("rij, sij -> rs", A, A)
Y = torch.einsum("rij, sij -> rs", B, B)
Z = torch.einsum("rij, sij -> rs", C, C)
torch.einsum("ij, ij, ij -> ", X, Y, Z)

which performs the following summation

$\sum_{r,s=1}^\sigma\left(\sum_{i,j=1}^\delta A^r_{ij} A^s_{ij} \right )\left(\sum_{k,l=1}^\delta B^r_{kl}B^s_{kl} \right )\left(\sum_{m,n=1}^\delta C^r_{mn}C^s_{mn} \right )$

Is it possible to formulate this in a more succinct (that is more vectorised/optimal) way? (e.g. using the fact that X, Y, Z are symmetric matrices)

user_2153636 · Answer

You can make it more succinct but I don’t see much room for actual performance optimisation:

X = torch.einsum("rij, sij", A, A)
Y = torch.einsum("rij, sij", B, B)
Z = torch.einsum("rij, sij", C, C)
torch.einsum("ij, ij, ij", X, Y, Z)

Simplify pytorch einsum

1 Answers