Einsum optimization

[ffuuugor]

A lot of our grad sample computations are based on einsum.
There's an easy (one-liner) way to optimize the expression: https://optimized-einsum.readthedocs.io/en/stable/

We want to try this and evaluate the impact

[ffuuugor]

Small update: now that we have benchmarks available, there's an easy way to evaluate the impact of this optimization.

See this guide on how to run benchmarks.