einops review
einops 0.8.2 gives tensor shape changes a named pattern: `b c h w -> b (c h w)` says which four axes arrive and exactly how three leave. Its functions rearrange, reduce, repeat, pack, unpack, and contract arrays while preserving the backend's tensor type. NumPy, PyTorch, JAX, TensorFlow, MLX, and several other frameworks are supported. Our install was 1 MB, had 0 direct dependencies, and included `py.typed`. This release adds the MLX backend, relies on native `torch.compile` support in PyTorch 2.8 or newer, and raises the minimum Python version to 3.9.
einops 0.8.2 installed in 0.2 seconds with 0 direct dependencies and used 1 MB in our sandbox, so tensor-heavy projects pay little for readable, cross-framework shape patterns. Native backend calls are easier to justify when a module has only one or two simple reshapes.
We installed it
| Install | ✓ · 0.2s | 1 package on disk · 1 MB |
| Import | ✓ | import einops in 0.10s · pure Python · py.typed · requires Python >=3.9 |
| Known vulns | 0 | (pip-audit) |
Answers from our run
Does einops install cleanly?
Yes. In a fresh container with an empty cache, pip install einops finished in 0.2s, leaving 1 package and 1 MB on disk. pip-audit reported no known vulnerabilities.
What does einops need to run?
Python >=3.9, and nothing compiled: it is pure Python. In our run import einops succeeded in 0.10s, and the package ships py.typed for type checkers.
einops or numpy: which should you use?
numpy: Use native reshape, transpose, repeat, and einsum when arrays stay in NumPy and the axis work remains short. einops 0.8.2 installed in 0.2 seconds with 0 direct dependencies and used 1 MB in our sandbox, so tensor-heavy projects pay little for readable, cross-framework shape patterns.
When should you not use einops?
The module contains one obvious reshape on a single backend. Native reshape and transpose calls avoid adding a pattern language for little gain.
Use it if
- Transformer code repeatedly splits attention heads, converts image grids to patch sequences, or joins token groups.
- Reviewers need the expected axis names and output order beside each reshape instead of reconstructing positional indexes.
- One tensor algorithm must keep the same shape notation across NumPy, PyTorch, JAX, TensorFlow, or MLX.
- A framework model needs rearrangement or reduction as a serializable layer inside Sequential-style composition.
- The module contains one obvious reshape on a single backend. Native `reshape` and `transpose` calls avoid adding a pattern language for little gain.
- Shape correctness must be expressed in function annotations before execution. jaxtyping covers annotated shapes and dtypes; einops validates a pattern only when the operation runs.
- The operation needs interpolation, convolution, sorting, masking, or arbitrary indexing. einops 0.8.2 handles axis arrangement, repetition, reduction, packing, and einsum-style contraction.
- Plain function calls must work under `torch.jit.script`. The README limits scripting support to the PyTorch layer classes, although function calls work with `torch.compile` on PyTorch 2.8 or newer.
- A hot loop handles tiny tensors and every wrapper call matters. The project documents readable shape operations, not a promise that pattern parsing beats a backend's direct primitive.
- Your team is unwilling to review grouped-axis order carefully. `(h patch)` and `(patch h)` can produce the same dimensions while arranging elements differently.
Setup reality
Our fresh Python 3.12 sandbox installed einops 0.8.2 in 0.2 seconds. The result was 1 package and 1 MB on disk, with 0 direct dependencies and no native compilation. import einops worked in 0.10 seconds. pip-audit found 0 known vulnerabilities. The package is pure Python, requires Python 3.9 or newer, uses the MIT license, and ships py.typed for type-checker discovery.
No credentials or config file are involved, but einops does not install a tensor framework. Bring NumPy, PyTorch, JAX, TensorFlow, MLX, or another documented backend yourself. Layer imports are framework-specific, such as einops.layers.torch.Rearrange. For implementations of the Python Array API, the project recommends functions from einops.array_api; version 0.8.2 also lets the main einops functions dispatch to MLX.
The first-run cost is learning the pattern rules. Parentheses join or split axes, repeat may introduce a new axis, and any input axis missing from a reduce result is reduced. A split dimension must factor exactly into the supplied lengths. Order inside a grouped axis is data layout, so changing (height patch) to (patch height) changes element placement even when both outputs report the same shape.
PyTorch 2.8 and later can compile einops function calls without the registration hook used by older combinations. The PyTorch layer classes support torch.jit.script; the plain functions do not. einops 0.8.2 leaves device placement, dtype, gradients, allocation, and kernel selection to the backend. Test the exact backend and compiler path used in production, especially when shapes vary between calls.
Patterns
Move time behind channels reorder-axes
from einops import rearrange
y = rearrange(x, 'time batch channels -> batch channels time')All 3 input axes appear in the output; a missing axis would require `reduce` instead of `rearrange`.
Flatten channel and spatial axes flatten-image
from einops import rearrange
y = rearrange(
x,
'batch channels height width -> batch (channels height width)',
)The left side asserts a 4-dimensional input before joining channels, height, and width.
Assert known axis lengths check-dimensions
y = rearrange(
x,
'batch channels height width -> batch (channels height width)',
channels=256,
height=19,
width=19,
)einops 0.8.2 raises when any supplied length disagrees with the runtime shape.
Restore tokens to an image grid restore-grid
grid = rearrange(
tokens,
'batch (height width) channels -> batch height width channels',
height=14,
width=14,
)The token axis must contain exactly 196 elements because 14 multiplied by 14 determines the split.
Convert images into patch tokens patchify-images
patches = rearrange(
images,
'batch channels (height p1) (width p2) -> batch (height width) (p1 p2 channels)',
p1=16,
p2=16,
)Both spatial dimensions must divide by the 16-element patch axes without a remainder.
Expose and merge attention heads split-attention-heads
q = rearrange(
q, 'batch tokens (heads dim) -> batch heads tokens dim', heads=8
)
out = rearrange(
q, 'batch heads tokens dim -> batch tokens (heads dim)'
)The feature width must divide by 8, and the grouped-axis order must match when the heads are merged again.
Average each spatial block average-pool
from einops import reduce
pooled = reduce(
x,
'batch channels (height h2) (width w2) -> batch channels height width',
'mean',
h2=2,
w2=2,
)The absent `h2` and `w2` axes are reduced, and both source dimensions must divide by 2.
Reduce both spatial axes global-max-pool
from einops import reduce
features = reduce(
x, 'batch channels height width -> batch channels', 'max'
)Version 0.8.2 applies `max` over the 2 named axes omitted from the result.
Copy grayscale data into three channels repeat-channels
from einops import repeat
rgb = repeat(
gray, 'height width -> height width channels', channels=3
)`repeat` can create the 3-value channel axis; `rearrange` cannot duplicate source values.
Compute attention scores with named axes attention-scores
from einops import einsum
scores = einsum(
q,
k,
'batch head query dim, batch head key dim -> batch head query key',
)einops places the equation after the 2 tensors and accepts multi-character axis names.
Join and restore different token groups pack-token-groups
from einops import pack, unpack
all_tokens, packed_shapes = pack(
[class_token, image_tokens, text_tokens], 'batch * channels'
)
encoded = transformer(all_tokens)
class_out, image_out, text_out = unpack(
encoded, packed_shapes, 'batch * channels'
)Keep the shape metadata returned by `pack`; `unpack` needs it to recover all 3 original groups.
Put a reshape inside a PyTorch model pytorch-layer
from torch.nn import Linear, Sequential
from einops.layers.torch import Rearrange
head = Sequential(
Rearrange('batch channels height width -> batch (channels height width)'),
Linear(16 * 5 * 5, 120),
)Import the framework-specific layer. The README says PyTorch layers support `torch.jit.script`, while plain operation functions do not.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| numpy | PyPI | Use native reshape, transpose, repeat, and einsum when arrays stay in NumPy and the axis work remains short. |
| torch | PyPI | Use view, reshape, permute, flatten, and native modules when a PyTorch-only team wants framework-specific code. |
| opt-einsum | PyPI | Use it when choosing and reusing efficient contraction paths matters more than named reshape patterns. |
| jaxtyping | PyPI | Use it when array shapes and dtypes belong in annotations with optional runtime checking. |
More ai / ml guides
openai · mcp · huggingface-hub · @modelcontextprotocol/sdk · scikit-learn · tiktoken · the whole shelf →
How this guide is made: grounded in the library's documentation, release notes, changelog, and issue history, on a fixed rubric — not a hands-on install of every release. The 50 most-downloaded entries are additionally install-verified in clean containers. Corrections: contact the desk.