einops: Einstein-Inspired Tensor Notation Across Five Frameworks
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
At a glance
- What is it?
- einops replaces ad hoc reshape, permute and repeat calls with a small pattern language that works on numpy, PyTorch, JAX, MLX and TensorFlow. The notation is the easy part; the harder questions are backend parity, layer serialization and when plain framework ops are the better call.
- Who is it for?
- Adopt einops if your codebase already mixes frameworks or if tensor axis bookkeeping is a recurring source of bugs in attention, patch embedding or multi-head code, because a single pattern string replaces a chain of view, permute and expand calls and the same string runs on numpy, torch, jax, mlx and tensorflow. Skip it if you need dynamic shapes decided at runtime, if you depend on a backend the project does not list, or if a one-line torch.permute already reads clearly.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 21 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The axis bookkeeping problem einops was written to remove
The library targets a specific kind of bug: tensor code where the meaning of an axis lives in a comment rather than in the expression. A patch embedding that has to be flattened, a multi-head attention block that splits a channel dimension into heads and later merges it back, a batch of images that has to be tiled along a new axis. In plain framework code each of those becomes a sequence of view, permute, reshape, expand and unsqueeze calls, and the reader has to reconstruct the shape arithmetic from the argument order. The README describes the project as providing "flexible and powerful tensor operations for readable and reliable code", and the reliability claim rests on the pattern being checked against the actual tensor shape at call time rather than inferred by the reader.
The intended audience is narrow but deep: people writing model code in numpy, pytorch, jax, mlx or tensorflow who move between those frameworks, and people who maintain layers that will be read by someone else six months later. It is not a numerical library. It does not add new math, and it does not replace autograd, kernels or memory planning. What it replaces is the transcription step between the shape you have in your head and the shape your framework wants.
Three functions, five verbs, and a pattern string that carries the shape contract
The core API is three functions imported from the package root: rearrange, reduce and repeat. Each takes a tensor, a pattern string and optional keyword arguments. The pattern has a left side describing input axes and a right side describing output axes, separated by an arrow. rearrange('t b c -> b c t') transposes. reduce('b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2) combines a rearrangement with a reduction, where the parenthesised groups on the left are split using the named sizes. repeat('h w -> h w c', c=3) copies along a new axis.
The mechanism is worth stating plainly because it explains both the ergonomics and the cost. The pattern is parsed, the left side is matched against the runtime shape, and the library derives the sequence of primitive operations needed to reach the right side. Axis names are arbitrary identifiers, not fixed letters, so 'batch time channel -> batch channel time' is as valid as 'b t c -> b c t'. That is the design decision that makes the code self-documenting: the name carries meaning, and the same name appearing twice in a pattern asserts that the two positions describe the same extent.
Two later additions extend the vocabulary. pack and unpack handle reversible packing of several tensors into one, and the README's example packs a class token, image tokens and text tokens of different dimensionality with the pattern 'b * c', then unpacks them after a transformer pass using the state object returned by pack. The asterisk absorbs the axes that differ between the inputs. einsum is the third verb, described as generic dot-product with three differences from the framework einsum: axes can be multi-lettered, the pattern goes last, and it works across frameworks. The README's example, einsum(A, B, 'b t1 head c, b t2 head c -> b head t1 t2'), is the attention score computation written with named axes instead of single letters.
Backend coverage is the real feature, and the release notes are where the gaps show
The README lists numpy, pytorch, jax, mlx and "others", and the repository topics add cupy, tensorflow and the array API standard. Release notes fill in the timeline: 0.7.0 added no-hassle torch.compile and array API standard support, 0.8.0 added a tinygrad backend, 0.8.2 added a full MLX backend and set Python 3.9 as the minimum, and 0.9.0dev is described as a development release for public testing, mostly typing changes.
That history is the useful part. It tells you the project treats backends as a first-class surface rather than a single-framework library with adapters bolted on, and it also tells you the edges are moving. A backend described as "full" in 0.8.2 was presumably partial before that. If you are on MLX, the version number matters more than the feature list on the homepage. If you are on a framework mentioned only in the topics list, the README does not tell you which operations are implemented there, and you should check the per-framework layer modules rather than assume parity with pytorch.
Layer wrappers follow the same split. The README shows Rearrange and Reduce imported from einops.layers.torch, einops.layers.tensorflow, einops.layers.flax and einops.layers.paddle, with a note that the code in other frameworks is almost identical. "Almost" is doing real work in that sentence: the module paths differ, and the serialization behaviour of a layer inside a Sequential or a Module will depend on the framework's own conventions, not on einops.
Getting it running, and the layer form that removes forward methods
Installation is one command, with a uv variant noted in the README:
pip install einops
From there the three core functions are imported directly:
from einops import rearrange, reduce, repeat
For model code the layer form is usually more useful, because it makes the pattern part of the module graph rather than a line inside forward:
from einops.layers.torch import Rearrange
The README's example puts Rearrange('b c h w -> b (c h w)') between a MaxPool2d and a Linear layer inside a torch.nn.Sequential, and notes that flattening no longer needs a hand-written forward. That is a small change with a real consequence: the shape transformation becomes visible in the model definition and, for frameworks that serialize module structure, it travels with the checkpoint. The README also states that operations and layers can be torch.compile'd, which is the 0.7.0 change showing up in user-facing form.
pack and unpack are imported from the same root:
from einops import pack, unpack
There is no config file, no environment variable and no registration step. The whole surface is the import, the pattern string and the keyword sizes. That is a point in its favour for review: a diff that changes a pattern is a diff a reviewer can read without knowing the surrounding shape arithmetic.
Where the notation stops helping
The pattern is a static description. Axes that are grouped on the left must be splittable by the sizes you supply, and axes that appear on the right must come from somewhere on the left or from a keyword. That is what makes the notation checkable, and it is also the boundary. If the split factor depends on a runtime value you do not want to thread through as a keyword, the pattern cannot express it, and you are back to framework ops. If your reshaping depends on data content rather than shape, einops has nothing to offer.
There is a second cost that the README does not discuss: the pattern string is not type-checked by your editor. The 0.9.0dev release note describes an overhaul of typing in einops, and the 0.8.2 note mentions relying on torch.compile, which suggests the project is aware that static analysis of these call sites is an open area. Until that typing work lands in a stable release, a misspelled axis name or a wrong size keyword is a runtime error, not a red squiggle. That is still better than silently wrong shape arithmetic, but it is not the same as compile-time safety.
Finally, einops is the wrong tool when the operation is already a single framework call with an obvious name. A one-line torch.permute in a function that does nothing else does not need a pattern string, and adding one introduces a parser, a dependency and a vocabulary the next reader has to learn. The library earns its place when the alternative is three or more chained shape operations, or when the same code has to run under more than one framework.
How it differs from framework-native reshape and permute
The obvious alternative is the framework's own tensor manipulation API: torch.view, torch.permute, torch.expand, torch.flatten, torch.einsum, and the numpy equivalents. The difference is not capability, since einops is documented as covering stacking, reshape, transposition, squeeze and unsqueeze, repeat, tile, concatenate, view and reductions, all of which the frameworks already do. The difference is where the shape contract lives.
With native ops, the contract is implicit in argument order and in the current shape of the tensor. With einops, it is written in the call. reduce('b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2) states that the channel axis is a product of h and h2 and that the width axis is a product of w and w2, and the library fails loudly if that is not true. The equivalent chain of view, permute and mean calls states the same thing across four lines and will happily produce a wrong result if the axis order was misremembered. That failure mode, a silent wrong answer rather than an exception, is the argument for the notation.
The second alternative is framework-native einsum. The README positions einops.einsum against it on three axes: multi-lettered axis names, pattern last, and cross-framework behaviour. The first is cosmetic but real for attention code where 'head' reads better than 'h'. The third is the substantive one, and it is the same argument as for the rest of the library: one pattern, several runtimes.
Maintenance, versioning and what the MIT licence does and does not cover
The project is active, not archived, with a last push in August 2026 and a development release in July 2026. The version history shows a steady cadence rather than a burst: 0.8.1 in February 2025, 0.8.2 in January 2026, 0.9.0dev in July 2026. The 0.8.2 notes bundle three unrelated changes (MLX backend, reliance on torch.compile, Python 3.9 minimum), which is worth knowing if you pin versions: a backend addition and a minimum-version bump arrived in the same release, so an upgrade can change both your supported Python floor and your compile behaviour.
The upgrade cost for most users is low because the public API is three functions plus two packing helpers plus einsum, and the README does not describe a deprecation cycle or breaking changes to the pattern syntax. The cost that does exist is in your own code: pattern strings are literals scattered through model definitions, and if a pattern's meaning changes when you refactor an axis, nothing in the type system will flag it. That is a grep-and-read cost, not a migration cost.
The licence is MIT, which permits commercial and closed-source use and requires only that the licence notice be preserved. That is a statement about the licence text, not legal advice; if your organisation has a policy on dependency licences, the MIT identifier is what to check against it. Note that the layer wrappers and EinMix are part of the same package under the same licence, so there is no separate commercial tier or contributor licence agreement to account for.
Editorial conclusion
Adopt einops if your codebase already mixes frameworks or if tensor axis bookkeeping is a recurring source of bugs in attention, patch embedding or multi-head code, because a single pattern string replaces a chain of view, permute and expand calls and the same string runs on numpy, torch, jax, mlx and tensorflow. Skip it if you need dynamic shapes decided at runtime, if you depend on a backend the project does not list, or if a one-line torch.permute already reads clearly. Verify three things before committing: that the backend you actually run matches the operation you need (pack and unpack, einsum or EinMix are not available everywhere), that the installed version's layers serialize correctly in your checkpoint format, and that your minimum Python version meets the 3.9 floor set in the 0.8.2 release notes. The MIT licence imposes no obligation on your source, but the pattern strings themselves are code, and they are the part you will maintain.
Community notes