einops

einops is a Python package that enables more readable tensor manipulation in existing libraries such as JAX, PyTorch etc. It does so by providing a minimalist API which uses Einstein-style notation. Some exa

Some side-by-side examples of PyTorch code with and without einops are shown here.