amazing at consolidating my code into something more readable
Readable if you're already familiar with einsum notation. Otherwise there's learning curve. An alternative to einsum is using multiple dot product and reshape ops, hopefully with each one having a comment - this would be a lot more readable imo.