Simplify pytorch einsum
Consider the following pytorch snippet:
X = torch.einsum(&quot;rij, sij -&gt; rs&quot;, A, A)
Y = torch.einsum(&quot;rij, sij -&gt; rs&quot;, B, B)
Z = torch.einsum(&quot;rij, sij -&gt; rs&quot;, C, C)
torch.einsum(&quot;ij, ij, ij -&gt; &quot;, X, Y, Z)

which performs the following summation

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)

You can make it more succinct but I don’t see much room for actual performance optimisation:
X = torch.einsum(&quot;rij, sij&quot;, A, A)
Y = torch.einsum(&quot;rij, sij&quot;, B, B)
Z = torch.einsum(&quot;rij, sij&quot;, C, C)
torch.einsum(&quot;ij, ij, ij&quot;, X, Y, Z)

