# einops: readable tensor reshaping for PyTorch, JAX, MLX and NumPy

> einops replaces reshape, permute and repeat calls with a notation string, and ships the same API across numpy, pytorch, jax, mlx and others. Here is what it does, how to install it, and where it stops being the right tool.

**arogozhnikov/einops** — Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

- Repository: https://github.com/arogozhnikov/einops
- Website: https://einops.rocks
- Stars: 9,605 · Forks: 402
- Language: Python
- License: MIT
- Published: 2026-09-10 · Updated: 2026-09-10 · Language: en
- Canonical page: https://hysenlabs.com/projects/arogozhnikov-einops

## The problem einops solves: tensor code that nobody can read

A typical deep learning forward pass is punctuated by lines like x.permute(0, 2, 1).reshape(b, -1) or x[:, None, :, None]. Each one is correct only if the reader remembers the axis order at that point in the function, and each one breaks silently when someone upstream changes the layout. The README frames the project as "flexible and powerful tensor operations for readable and reliable code", and the reliability part is the interesting claim: the pattern string is checked against the actual tensor shape at runtime, so a mismatch raises instead of quietly producing a wrong-shaped result.

The audience is anyone who moves axes around in Python. The README lists numpy, pytorch, jax, mlx and "others", and the repository topics also name cupy and tensorflow. That breadth is the point. A pattern written for a pytorch tensor is the same string you would write for a numpy array, so a function can be ported between frameworks without rewriting its index arithmetic. The project is not a new tensor library and does not want to be one. It sits on top of whatever array type you already have, which is why pyproject.toml declares no run-time or installation-time dependencies at all.

## How einops works: patterns, parsing and a per-framework dispatch

Every einops call takes a tensor, a pattern string, and sometimes arguments. The pattern has a left side describing the input axes and a right side describing the output axes, separated by an arrow. Axis names are arbitrary identifiers, and the parentheses carry the structural information: axes inside one pair of parentheses are flattened together on that side of the arrow, and axes that appear on one side but not the other are created or removed.

That single rule covers a surprising amount of ground. The README states that the three core operations cover "stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions". Rearrangement is one family, reduction is another (it takes an operation such as 'mean' plus the axes being reduced, given as keyword arguments), and repetition is a third. Because the pattern is parsed rather than executed as a chain of calls, einops can validate that the axis names and the tensor rank agree before any work happens.

The repository layout reflects the multi-backend design. The einops/ package holds the core, and einops.layers holds framework-specific layer wrappers, with the README showing imports from einops.layers.torch, einops.layers.tensorflow, einops.layers.flax and einops.layers.paddle. The wheel build in pyproject.toml ships exactly two packages, einops and einops.layers, so the framework adapters live inside the same distribution rather than as separate installs. The release notes for 0.8.2 mention a full MLX backend and a move to rely on torch.compile; 0.8.0 added a tinygrad backend. The current development release, 0.9.0dev, is described in its own release note as "mostly typing changes".

## Installing einops and running a first rearrange

The README gives a single install line, and notes that uv works too. There is no compiled extension and no dependency resolution to worry about, because the dependency list in pyproject.toml is empty.

```bash
pip install einops
```

After that, the three core functions import from the top-level package. The README's micro-reference shows one line each for rearrangement, reduction and repetition:

```python
from einops import rearrange, reduce, repeat
output_tensor = rearrange(input_tensor, 't b c -> b c t')
output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)
```

The first call is a plain axis permutation. The second is the more interesting one: the input axes c and w are each split into two named factors, h2 and w2 are supplied as keyword arguments, and the reduction averages over them while leaving the rest of the layout intact. That is a patch-merge operation in one line. The third call adds a new axis of size c and copies the input along it.

For use inside a model rather than in a script, the layers subpackage gives you module wrappers. The README's pytorch example replaces a flatten step with a Rearrange layer:

```python
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),
)
```

The README notes that operations and layers can be torch.compile'd, and that torch.jit.script is supported for the pytorch layers. The project also publishes a browser playground that the README says can run two of the four example notebooks without a local install, which is a reasonable way to try the notation before adding a dependency.

## pack and unpack, and where einsum fits

Two later additions are worth separating from the core three. pack and unpack handle the case where several tensors of different rank need to travel through one function and come back out. The README's example packs a class token, image tokens and text tokens into a single tensor with the pattern 'b * c', where the asterisk absorbs the differing middle axes, and returns a list of lengths that unpack uses to reverse the operation.

```python
from einops import pack, unpack
packed, ps = pack([class_token_bc, image_tokens_bhwc, text_tokens_btc], 'b * c')
class_emb_bc, image_emb_bhwc, text_emb_btc = unpack(transformer(packed), ps, 'b * c')
```

The README argues this is "better than stack/split/concatenate", and the practical difference is that the bookkeeping of how long each segment was is returned to you rather than recomputed at the call site. If you have ever written a transformer that concatenates a class token with patch embeddings and then slices the result back apart, this is the operation being replaced.

einsum is the other addition, introduced in 0.5 according to the update list. It differs from the einsum you may already know in three stated ways: axes can be multi-lettered, the pattern goes last rather than first, and it works across the supported frameworks. The README example contracts two attention-shaped tensors:

```python
from einops import einsum
C = einsum(A, B, 'b t1 head c, b t2 head c -> b head t1 t2')
```

Multi-lettered axis names matter more than they sound. In a raw einsum call, 'b t1 head c' would have to be written as single characters, and once you have more than a handful of axes the letters stop meaning anything. Naming an axis 'head' keeps the pattern legible six months later.

## Where einops is the wrong choice

The first limitation is the one the project itself is quiet about: einops validates and dispatches, so it is an extra layer over the backend's native operations. The README does not publish any timing comparison, and the search questions people ask about whether einops is slow or fast are not answered in the README text. If your workload is dominated by one heavily optimized kernel, wrap that kernel directly and keep einops at the boundaries where shapes change.

The second is scope. einops rearranges, reduces and repeats. It does not manage device placement, dtypes, gradient flow or memory layout, and it is not a replacement for the array library you are already using. A pattern like 'b c h w -> b (c h w)' tells you the shape of the result and nothing about whether the result is contiguous or whether the copy is avoidable.

The third is the typing story. The 0.9.0dev release is explicitly a development release for public testing, and its release note says the work is mostly typing changes. If your team relies on static type checkers or on the torch.jit.script support mentioned in the README, the stable line to pin is 0.8.2. Adopting a dev release to get better annotations is a reasonable trade only if you are prepared to track the API as it settles.

Finally, the minimum Python version in pyproject.toml is 3.10. The 0.8.2 release note mentions setting Python 3.9 as the minimum, so the requirement has moved forward since. Environments pinned to older interpreters cannot use the current package.

## einops compared with raw einsum and raw framework calls

The nearest alternative for contraction is the einsum that already ships with numpy and pytorch. It solves an overlapping problem with a different notation: single-character axis labels, the pattern placed before the operands, and no support for the rearrange, reduce and repeat families. If all you need is a batched matmul with an unusual axis order, calling torch.einsum directly avoids a dependency and keeps the operation inside the framework you already trust. The trade is legibility on tensors with many axes, and portability if the same code has to run on a second backend.

The other alternative is doing nothing and keeping the permute and reshape chain. That is genuinely fine for two or three operations in a small file. It stops being fine when the same function has five of them, when the axis order is not obvious from the variable names, or when the code has to run under more than one framework. The README's list of covered operations (stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and reductions) is a fair description of the surface area you would otherwise hand-write.

A third comparison point is the array API standard. The 0.7.0 release note mentions support for it, which means einops is not inventing a private abstraction so much as layering a notation on top of an interface that several frameworks have agreed on. That is a meaningful difference from a library that only works with one vendor's tensors.

## Maintenance, licensing and what upgrading costs

The repository is not archived, and the last push was on 2026-08-26. The release cadence visible in the notes is not rapid: 0.8.1 in February 2025, 0.8.2 in January 2026, and a 0.9.0dev development release in July 2026. That pattern suggests a project that ships when there is something to ship rather than on a schedule, which matters if you are waiting on a specific fix.

The upgrade cost is low by construction. The API is three core functions plus pack, unpack and einsum, and the pattern strings are data rather than code, so a version bump does not rewrite your call sites. The main compatibility lever is the Python floor, which pyproject.toml sets at 3.10. The typing overhaul in 0.9.0dev is the one change that could surface errors in code that passes type checking today, which is an argument for pinning 0.8.2 in CI while trying 0.9.0dev separately.

Licensing is MIT, declared both in the repository LICENSE file and as license = { text = 'MIT' } in pyproject.toml. MIT is permissive and imposes no copyleft obligation on your own code. That is a statement about the licence text, not legal advice; if your organisation has a policy on dependency licences, route it through whoever normally handles that. The practical point is that there is no dual-licensing scheme or commercial tier mentioned anywhere in the README or pyproject.toml, so there is nothing to negotiate.

## Conclusion

Adopt einops if your codebase already does a lot of permute, reshape, squeeze and repeat on tensors in numpy, pytorch, jax, mlx or the other supported backends, and you want those operations stated as a pattern instead of a chain of method calls. Do not adopt it as a replacement for einsum-style contraction in a hot loop without checking the backend notes, and do not expect it to manage device placement, dtypes or autograd for you. Before committing, verify three things: that your framework appears in the supported list for the version you install, that the minimum Python requirement of 3.10 from pyproject.toml is acceptable in your environment, and that your existing reshape-heavy code produces the same output after conversion. The typing overhaul announced for 0.9.0 is still a dev release, so pin to 0.8.2 if you depend on stable type stubs.

## FAQ

### Does einops work with JAX?

Yes. The README lists jax among the supported frameworks and the repository topics include jax. Patterns are written the same way as for numpy or pytorch, and there is a separate einops.layers.flax module for layer wrappers.

### How do I use einops.rearrange in Python?

Import rearrange from einops, then pass a tensor and a pattern with an arrow between the input and output axes, for example rearrange(input_tensor, 't b c -> b c t'). Axes inside parentheses are flattened together on that side of the arrow.

### How do I install einops?

The README gives a single command, pip install einops, and notes that uv pip install einops works as well. The package declares no run-time or installation-time dependencies in pyproject.toml.

### What is einops?

einops stands for Einstein-Inspired Notation. It is a Python library that expresses tensor rearrangement, reduction and repetition as pattern strings, and it supports numpy, pytorch, jax, mlx and other frameworks from the same API.

### Is einops slow?

The README does not publish any timing or benchmark comparison, so no performance claim can be made here. einops parses and validates a pattern before dispatching to the underlying framework, which is an extra layer over calling that framework directly.

### What is the difference between einops and einsum?

einops provides einsum alongside rearrange, reduce, repeat, pack and unpack. Its einsum differs in three stated ways: axes can be multi-lettered, the pattern goes last instead of first, and it works across the supported frameworks rather than one.

## Sources

- [arogozhnikov/einops on GitHub](https://github.com/arogozhnikov/einops)
- [License: MIT](https://github.com/arogozhnikov/einops/blob/main/LICENSE)
- [Project website](https://einops.rocks)
- [README](https://github.com/arogozhnikov/einops/blob/main/README.md)
- [Releases](https://github.com/arogozhnikov/einops/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/arogozhnikov-einops