About this project
einops is an open-source Python library that provides flexible and powerful tensor operations through an Einstein-inspired notation, designed to make tensor manipulation code more readable and reliable. The name stands for Einstein-Inspired Notation for operations, loosely inspired by numpy.einsum.
The core API is minimalistic yet powerful, built around three main functions: rearrange (rearrange elements according to a pattern), reduce (combine rearrangement with reduction such as mean or max), and repeat (copy along a new axis). Later additions include pack and unpack for reversibly combining multiple tensors of different dimensionality, and einsum with support for multi-lettered axis names. Einops also ships an EinMix layer, a generic linear layer well suited for MLP Mixers and similar architectures, plus framework-specific layer wrappers (Rearrange, Reduce, etc.) for PyTorch, TensorFlow, Flax, and Paddle.
The notation emphasizes semantic information: instead of writing x.view(x.shape[0], -1), users write rearrange(x, 'b c h w -> b (c h w)'), which documents both input and output shapes. Named axes like 'time c h w -> time (c h w)' hint that the data is a video sequence rather than an independent batch. Users can also pin specific dimensions (e.g. c=256, h=19, w=19) so the operation validates shapes at runtime rather than relying on comments. Results are strictly determined, so ambiguous operations like depth-to-space can be written unambiguously. The same syntax works uniformly across 1D, 2D, and 3D pooling, and behavior is consistent across frameworks even where native APIs differ (for example, flatten semantics vary between NumPy/PyTorch and TensorFlow/Keras).
Einops supports NumPy, PyTorch, TensorFlow, JAX, CuPy, Flax, Paddle, OneFlow, tinygrad, and PyTensor, and works with any framework implementing the Python array API standard, including NumPy >= 2.0, MLX, pydata/sparse, cubed, and nndonnx. It supports torch.compile and torch.jit.script for PyTorch layers. The library requires Python 3.10 or later and is installed via pip install einops. The project was presented as an oral at ICLR 2022 and is used by more than 10,000 GitHub projects.
Comments
0 Rating appears after 10 ratings
Sign in to join the discussion.