Comment by mcemilg

Looks wonderful. But I would like to add this, I hate einops, it doesn't make it simple to read unfortunately.