einops
einops gives you a handful of functions (rearrange, reduce, repeat, plus pack, unpack and einsum) that express tensor shape operations as readable patterns like 'b c h w -> b (c h w)'. One notation replaces reshape, transpose, squeeze, stack, tile and their framework-specific quirks, works identically across numpy, PyTorch, JAX, TensorFlow, MLX and anything implementing the array API standard, and validates dimensions at runtime so shape bugs fail loudly with named axes instead of silently producing a wrong tensor.
One of the rare dependencies that makes research code strictly more readable at near-zero runtime and install cost, which is why it appears in most modern transformer implementations. Skip it only if your team refuses the notation.
Use it if
- You write transformer or vision code with heavy dimension juggling (splitting attention heads, patchifying images, merging batch axes) and want each operation to document its own input and output shapes
- You want runtime shape checks: patterns assert the number of dimensions and any named sizes you pin down, so a wrong input errors immediately instead of propagating garbage
- You maintain code that must run across frameworks: the same rearrange call behaves identically in numpy, torch and jax, unlike flatten or repeat, which differ between them
- You build models with nn.Sequential and want Rearrange/Reduce layers so a flatten or pooling step does not require a custom forward method
- Your team will not learn the notation: it is another mini-language, and a project doing an occasional .reshape() gains nothing from the dependency or the review burden
- You expect more than shape manipulation: there is no math beyond reductions and einsum, so it complements your framework rather than shrinking it
- You need torch.jit.script: einops function calls work with torch.compile but are not scriptable, and only the torch layer classes survive scripting
- You are optimizing sub-microsecond hot paths on tiny tensors: pattern parsing is cached, but wrapper overhead versus the raw native op is still measurable in tight loops
Setup reality
pip install einops is a small pure-Python install with zero dependencies; no framework is pulled in because it dispatches to whichever backend your tensors already come from. The real cost is human: the notation takes an afternoon to internalize and everyone reviewing the code needs it too. Mechanical gotchas are few but real: layer classes live in per-framework modules (einops.layers.torch versus einops.layers.flax), array-API backends go through einops.array_api, torch.compile is fine while torch.jit.script is not for function calls, and recent releases require Python 3.10 or later.
Patterns
Reorder axes readablytranspose-axes
from einops import rearrange
# (time, batch, channels) -> (batch, channels, time)
y = rearrange(x, 't b c -> b c t')Every axis on the left must appear on the right (or be reduced away with reduce); a missing name is an immediate error, not a silent drop.
Flatten trailing dimensionsflatten-batch
# (batch, channels, height, width) -> (batch, features)
y = rearrange(x, 'b c h w -> b (c h w)')
# pin sizes to get a hard runtime check
y = rearrange(x, 'b c h w -> b (c h w)', c=256, h=19, w=19)Pinning sizes turns the pattern into an assertion; mismatched inputs raise with the axis name instead of reshaping into nonsense.
Split one axis into severalsplit-axis
# flat tokens back to a grid: (b, h*w, c) -> (b, h, w, c)
y = rearrange(x, 'b (h w) c -> b h w c', h=14, w=14)When decomposing an axis you must specify enough sizes for einops to solve the rest; giving only h works if w is implied by the total.
Pooling as a reduction patternspatial-pooling
from einops import reduce
# 2x2 average pooling
y = reduce(x, 'b c (h h2) (w w2) -> b c h w', 'mean', h2=2, w2=2)
# global max pool
y = reduce(x, 'b c h w -> b c', 'max')Reductions accept mean, max, min, sum and prod; the same pattern style covers 1d/2d/3d pooling without separate ops.
Broadcast or tile without tile/repeat confusionrepeat-along-new-axis
from einops import repeat
# grayscale (h, w) -> rgb (h, w, 3)
rgb = repeat(image, 'h w -> h w c', c=3)
# tile along width
wide = repeat(image, 'h w -> h (tile w)', tile=2)numpy tile and torch repeat mean different things; einops repeat is one spelling across all frameworks and the pattern shows which axis grows.
Patchify images for a ViTimage-to-patches
# (b, c, H, W) -> (b, num_patches, patch_dim)
patches = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=16, p2=16)H and W must be divisible by the patch size or the pattern raises; the error message names the failing axis, which beats debugging a reshape.
Split and merge attention headssplit-attention-heads
# (b, tokens, heads*dim) -> (b, heads, tokens, dim)
q = rearrange(q, 'b t (h d) -> b h t d', h=8)
# back after attention
out = rearrange(out, 'b h t d -> b t (h d)')Order inside parentheses matters: (h d) and (d h) index memory differently, and einops does exactly what the pattern says rather than guessing.
Attention scores with named-axis einsumeinsum-attention-scores
from einops import einsum
scores = einsum(q, k, 'b h t1 d, b h t2 d -> b h t1 t2')Unlike framework einsum, axes can be whole words and the pattern comes last; it dispatches to the backend's einsum underneath.
Pack mixed tensors into one sequence and unpack laterpack-unpack-tokens
from einops import pack, unpack
packed, ps = pack([cls_token, image_tokens, text_tokens], 'b * c')
out = transformer(packed)
cls_out, img_out, txt_out = unpack(out, ps, 'b * c')pack flattens whatever sits at the * position and remembers the shapes in ps, so unpack restores each tensor's original dimensionality exactly.
Use Rearrange as a layer in nn.Sequentialrearrange-layer-in-model
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrange
model = Sequential(
Conv2d(6, 16, kernel_size=5),
MaxPool2d(kernel_size=2),
Rearrange('b c h w -> b (c h w)'),
Linear(16 * 5 * 5, 120),
ReLU(),
)Import the layer from the framework-specific module (einops.layers.torch here); the torch layers are scriptable and compile-friendly.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| torch | PyPI | Native view/permute/flatten idioms when the project is PyTorch-only and the team prefers standard calls |
| opt-einsum | PyPI | Optimized contraction paths when your bottleneck is large einsum expressions rather than readability |
| jaxtyping | PyPI | Shape and dtype checking through type annotations when you want shape safety without changing runtime code |