Comment by machinekob

I would say https://github.com/arogozhnikov/einops should be pretty close to your requirments