There is a very nice library called Einops that makes it easier to work with einsum. See:

Descriptive Matrix Operations with Einops — https://www.kolaayonrinde.com/blog/2024/01/08/einops.html

Einops tutorial — https://einops.rocks/1-einops-basics/