Comment by _rerr

This is neat. However I've rather use einsum and einops [0] to reshape stuff in practice.

[0] https://github.com/arogozhnikov/einops