einops
A library for readable, flexible tensor operations, such as reshaping, transposition and reduction, using one consistent notation across NumPy, PyTorch, JAX, TensorFlow and other frameworks.
einops is a library for tensor operations that aims to make code more readable and reliable across frameworks, including NumPy, PyTorch, JAX, TensorFlow, CuPy, MLX and others. Three core operations, rearrange, reduce and repeat, cover stacking, reshaping, transposition, squeezing, tiling, concatenation, views and many kinds of reduction, each expressed as a short pattern string rather than a chain of separate calls.
EinMix extends the same pattern-based approach to mixing operations, and einsum notation is available directly through einops. It supports torch.compile and torch.jit.script for PyTorch layers, the array API standard, and packing and unpacking operations for combining and splitting tensors. An interactive playground can run some of its tutorial notebooks directly in a browser.
It is written in Python, installs as a package, and is released under the MIT license.