There is a very nice library called Einops that makes it easier to work with einsum. See:<p>Descriptive Matrix Operations with Einops — <a href="https://www.kolaayonrinde.com/blog/2024/01/08/einops.html" rel="nofollow">https://www.kolaayonrinde.com/blog/2024/01/08/einops.html</a><p>Einops tutorial — <a href="https://einops.rocks/1-einops-basics/" rel="nofollow">https://einops.rocks/1-einops-basics/</a>