[machinekob]

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

[polygamous_bat]

I would like to second that! I converted multiple of my lab mates into Einops just by having them browse through the tutorial in https://einops.rocks/ :)

[jpeanuts]

Looking at the documentation - einops seems to only implement operations on a single tensor, so it's quite far from a general replacement for einsum, which can perform tensor-products on an arbitrary number of tensors.