einops

einops gives NumPy, PyTorch, TensorFlow, JAX, and other array libraries a single, explicit notation for tensor operations that are normally expressed with error-prone chains of reshape, transpose, squeeze, and framework-specific tile/repeat calls. Its core rearrange, reduce, repeat, pack, and unpack functions take a plain-text pattern describing input and output axes, so the intent of a tensor operation is visible directly in the code instead of buried in index arithmetic.

Because the same pattern-based API works identically across supported backends (NumPy, PyTorch, TensorFlow, JAX, CuPy, PaddlePaddle, and any framework implementing the Python array API standard), einops code is portable between frameworks without behavioral surprises — e.g. flatten() producing different shapes in different libraries. It also ships framework-specific layers (Rearrange, Reduce, EinMix) for dropping the same notation directly into PyTorch, TensorFlow, or Flax model definitions.

What You Get

• rearrange, reduce, and repeat functions covering stacking, reshaping, transposition, squeeze/unsqueeze, tiling, and reductions with one consistent pattern syntax
• pack/unpack for reversibly combining tensors of different shapes into one, replacing manual stack/split/concatenate logic
• An einsum wrapper supporting multi-lettered axis names with the pattern specified last, unlike NumPy’s positional einsum
• Framework-specific layers (Rearrange, Reduce, EinMix) for PyTorch, TensorFlow, Flax, and PaddlePaddle model definitions
• Support for any framework implementing the Python array API standard (NumPy 2.0+, MLX, sparse, cubed) via einops.array_api

Common Use Cases

• Replacing chains of .view()/.transpose()/.reshape() calls in a PyTorch model with a single, self-documenting rearrange pattern
• Writing tensor-manipulation code once that runs unchanged across NumPy, PyTorch, JAX, and TensorFlow backends
• Building MLP-Mixer-style architectures with EinMix, a generic linear layer expressed in einops notation
• Combining and splitting multi-modal tensors (e.g. class tokens, image patches, text tokens) reversibly with pack/unpack

Under The Hood

Architecture - The core lives in einops/einops.py (the rearrange/reduce/repeat implementations), einops/parsing.py (parses the text pattern into an axis-transformation plan), and einops/packing.py (pack/unpack). einops/_backends.py implements a backend-abstraction layer that dispatches the parsed operation to the appropriate framework (NumPy, PyTorch, TensorFlow, JAX, CuPy, etc.), while einops/layers/ wraps the same operations as framework-native nn.Module/keras.Layer-style classes and einops/array_api.py targets the Python array API standard directly.

Tech Stack - Pure Python (69% of the repo, with the rest being Jupyter notebook tutorials) with zero runtime dependencies — einops ships no required dependencies at all, relying purely on whichever array framework the caller already has installed. Built with hatchling; framework-specific code paths (_torch_specific.py) are isolated so the core package doesn’t hard-depend on any single ML library.

Code Quality - Tests live under einops/tests/ (10 files) and are run via GitHub Actions across supported Python versions and backends. The library is narrowly scoped (a handful of core files) but heavily exercised in production — GitHub reports over 10,000 dependent projects — and has been stable enough that its pattern syntax hasn’t changed since early releases.

API Design - The pattern-string API ('b c h w -> b (c h w)') is deliberately declarative: it states what the input and output axes mean rather than how the reshape is computed, which also lets einops validate assumptions (e.g. c=256, h=19, w=19) that a bare .view() call cannot. This is a genuinely low learning curve for anyone who reads the README’s side-by-side comparisons, though newcomers unfamiliar with einsum-style notation need a short ramp-up before the patterns become second nature.