einops (arogozhnikov/einops)

einops

https://github.com/arogozhnikov/einops
Admin
einops provides flexible and powerful tensor operations for readable and reliable code, with support...

Tokens:67,452
Snippets:605
Trust Score:9.4
Update:1 week ago
Show doc for...
Context Summary (auto-generated)
Raw
# einops

einops is a Python library providing flexible and powerful tensor operations using Einstein-inspired notation for readable and reliable code. It works uniformly across all major deep learning and numerical computing frameworks — including NumPy, PyTorch, JAX, TensorFlow, MLX, CuPy, Flax, Paddle, OneFlow, and any library implementing the Python Array API standard — offering a single, consistent API regardless of the underlying framework.

The core of einops is a pattern language that explicitly names every axis of a tensor in an input→output transformation. This notation simultaneously documents intent, validates shapes at runtime, and expresses complex reshapes, transpositions, reductions, and repetitions in one readable expression. The library ships three fundamental functions (`rearrange`, `reduce`, `repeat`), two packing utilities (`pack`, `unpack`), a named-axis `einsum`, and neural-network layer wrappers (`Rearrange`, `Reduce`, `EinMix`) for PyTorch, TensorFlow, Flax, and Paddle.

## Installation

```bash
pip install einops
```

---

## `rearrange` — Reshape, Transpose, Squeeze, Unsqueeze, Stack, Concatenate

`rearrange` is the fundamental operation for reordering and reshaping tensor axes. It replaces `transpose`, `reshape`/`view`, `squeeze`, `unsqueeze`, `stack`, and `concatenate` with a single pattern-based call. The left side of the pattern names the input axes and the right side describes the desired output layout; parentheses combine or split axes.

```python
import numpy as np
from einops import rearrange

# --- Basic transpose: (batch, height, width, channel) -> (batch, channel, height, width) ---
x = np.random.randn(8, 32, 32, 3)
out = rearrange(x, 'b h w c -> b c h w')
assert out.shape == (8, 3, 32, 32)

# --- Flatten spatial dims into a single vector ---
flat = rearrange(x, 'b h w c -> b (c h w)')
assert flat.shape == (8, 3072)  # 3*32*32 = 3072

# --- Split an axis: decompose height into tiles ---
tiles = rearrange(x, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2)
assert tiles.shape == (32, 16, 16, 3)  # 8*2*2 batches of 16×16 tiles

# --- Stack a list of arrays along a new batch axis ---
images = [np.random.randn(30, 40, 3) for _ in range(16)]
batch = rearrange(images, 'b h w c -> b h w c')
assert batch.shape == (16, 30, 40, 3)

# --- Concatenate images horizontally; 640 = 16 * 40 ---
mosaic = rearrange(images, 'b h w c -> h (b w) c')
assert mosaic.shape == (30, 640, 3)

# --- Space-to-depth (pixel shuffle inverse) ---
s2d = rearrange(x, 'b (h h2) (w w2) c -> b h w (c h2 w2)', h2=2, w2=2)
assert s2d.shape == (8, 16, 16, 12)

# --- Squeeze / unsqueeze by using size-1 axes ---
squeezed = rearrange(np.zeros((1, 5, 1, 7)), '1 a 1 b -> a b')
assert squeezed.shape == (5, 7)

# --- Enforce concrete axis sizes as runtime checks ---
checked = rearrange(x, 'b h w c -> b (h w) c', h=32, w=32)
assert checked.shape == (8, 1024, 3)
```

---

## `reduce` — Rearrange Combined with Reduction

`reduce` extends `rearrange` with an explicit reduction step: axes present on the left but absent on the right are reduced using the specified operation. Supported built-in reductions are `'min'`, `'max'`, `'sum'`, `'mean'`, `'prod'`, `'any'`, `'all'`. A custom callable `f(tensor, axes: tuple) -> tensor` is also accepted.

```python
import numpy as np
from einops import reduce

x = np.random.randn(10, 3, 64, 64)  # batch=10, channels=3, h=64, w=64

# --- Global average pooling ---
gap = reduce(x, 'b c h w -> b c', 'mean')
assert gap.shape == (10, 3)

# --- 2-D max-pooling with 2×2 kernel ---
pooled = reduce(x, 'b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2)
assert pooled.shape == (10, 3, 32, 32)

# --- Anonymous kernel size syntax (equivalent to above) ---
pooled2 = reduce(x, 'b c (h 2) (w 2) -> b c h w', 'max')
assert pooled2.shape == (10, 3, 32, 32)

# --- Adaptive pooling to fixed output grid 4×4 ---
adaptive = reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h1=4, w1=4)
assert adaptive.shape == (10, 3, 4, 4)

# --- Subtract per-batch, per-channel mean (batch norm style) ---
normalized = x - reduce(x, 'b c h w -> b c 1 1', 'mean')

# --- Reduce over the batch dimension ---
batch_max = reduce(x, 'b c h w -> c h w', 'max')
assert batch_max.shape == (3, 64, 64)

# --- Uniform 1-D, 2-D, 3-D pooling with the same pattern style ---
x3d = np.random.randn(2, 8, 16, 16, 16)
pool3d = reduce(x3d, 'b c (x dx) (y dy) (z dz) -> b c x y z', 'max', dx=2, dy=2, dz=2)
assert pool3d.shape == (2, 8, 8, 8, 8)

# --- Custom reduction: log-sum-exp ---
import numpy as np
logsumexp = reduce(x, 'b c h w -> b c', lambda t, axes: np.log(np.sum(np.exp(t), axis=axes)))
assert logsumexp.shape == (10, 3)
```

---

## `repeat` — Tile and Broadcast Along New or Existing Axes

`repeat` adds axes (or enlarges existing ones) by repeating data. It subsumes NumPy's `tile`, PyTorch's `repeat`, and `broadcast_to`, with identical syntax across all frameworks.

```python
import numpy as np
from einops import repeat, reduce

image = np.random.randn(30, 40)  # grayscale H×W

# --- Grayscale → RGB by repeating along a new channel axis ---
rgb = repeat(image, 'h w -> h w c', c=3)
assert rgb.shape == (30, 40, 3)

# --- Add a batch dimension ---
batched = repeat(image, 'h w -> b h w', b=8)
assert batched.shape == (8, 30, 40)

# --- Tile: duplicate along height ---
tiled_h = repeat(image, 'h w -> (tile h) w', tile=2)
assert tiled_h.shape == (60, 40)

# --- Upsample 2× by repeating each pixel ---
up2x = repeat(image, 'h w -> (h h2) (w w2)', h2=2, w2=2)
assert up2x.shape == (60, 80)

# --- 'Pixelate': downsample then upsample ---
small = reduce(image, '(h h2) (w w2) -> h w', 'mean', h2=2, w2=2)
pixelated = repeat(small, 'h w -> (h h2) (w w2)', h2=2, w2=2)
assert pixelated.shape == image.shape

# --- Same call works in PyTorch, JAX, TF without any change ---
# import torch; image_t = torch.from_numpy(image)
# repeat(image_t, 'h w -> h w c', c=3)  # identical API
```

---

## `einsum` — Named-Axis Einsum Across Frameworks

`einops.einsum` is a drop-in improvement over framework-native einsum functions. It accepts tensors first and the pattern last, supports **multi-letter axis names** (making patterns self-documenting), and works identically across NumPy, PyTorch, JAX, and TensorFlow.

```python
import numpy as np
from einops import einsum

# --- Matrix multiplication with descriptive names ---
A = np.random.randn(4, 8, 16)   # batch, seq, features
B = np.random.randn(4, 16, 32)  # batch, features, out
C = einsum(A, B, 'batch seq features, batch features out -> batch seq out')
assert C.shape == (4, 8, 32)

# --- Attention score matrix (batch, heads, time × time) ---
Q = np.random.randn(2, 8, 64, 32)  # batch, heads, t_query, d_head
K = np.random.randn(2, 8, 64, 32)  # batch, heads, t_key,   d_head
scores = einsum(Q, K, 'b head t1 d, b head t2 d -> b head t1 t2')
assert scores.shape == (2, 8, 64, 64)

# --- Batched image filtering ---
images = np.random.randn(128, 16, 16)
filters = np.random.randn(16, 16, 30)
out = einsum(images, filters, 'batch h w, h w channel -> batch channel')
assert out.shape == (128, 30)

# --- Batched linear projection with ellipsis ---
data    = np.random.randn(50, 30, 20)  # ...batch, in_dim
weights = np.random.randn(10, 20)       # out_dim, in_dim
out = einsum(weights, data, 'out_dim in_dim, ... in_dim -> ... out_dim')
assert out.shape == (50, 30, 10)

# --- Matrix trace (scalar output) ---
M = np.random.randn(10, 10)
trace = einsum(M, 'i i ->')
assert trace.shape == ()
```

---

## `pack` and `unpack` — Reversible Multi-Tensor Packing

`pack` concatenates tensors of **different dimensionalities** into a single tensor by collapsing variable middle axes into one. The `*` wildcard in the pattern marks the "packed" region. `unpack` is the exact inverse given the packed shapes (`ps`) returned by `pack`. Together they replace `stack`/`concatenate`/`split` in multi-modal or multi-scale pipelines.

```python
import numpy as np
from einops import pack, unpack

# --- Pack tokens of different shapes for a transformer ---
class_token  = np.random.randn(2, 512)         # b c
image_tokens = np.random.randn(2, 16, 16, 512) # b h w c
text_tokens  = np.random.randn(2, 32, 512)     # b t c

packed, ps = pack([class_token, image_tokens, text_tokens], 'b * c')
# packed.shape == (2, 1+256+32, 512)  => all tokens in one sequence
print(packed.shape, ps)
# (2, 289, 512)  [(), (16, 16), (32,)]

# Restore original tensors after processing (e.g. transformer forward pass)
cls_out, img_out, txt_out = unpack(packed, ps, 'b * c')
assert cls_out.shape  == (2, 512)
assert img_out.shape  == (2, 16, 16, 512)
assert txt_out.shape  == (2, 32, 512)

# --- Mixed-dimensionality packing ---
from numpy import zeros as Z
inputs = [Z([2, 3, 5]), Z([2, 3, 7, 5]), Z([2, 3, 7, 9, 5])]
packed2, ps2 = pack(inputs, 'i j * k')
assert packed2.shape == (2, 3, 71, 5)   # 1 + 7 + 63 = 71
assert ps2 == [(), (7,), (7, 9)]

restored = unpack(packed2, ps2, 'i j * k')
assert [x.shape for x in restored] == [(2,3,5), (2,3,7,5), (2,3,7,9,5)]
```

---

## `parse_shape` — Named Shape Extraction and Validation

`parse_shape` maps axis names to their sizes for a given tensor, optionally skipping axes with `_`. The resulting dict can be passed directly as `**kwargs` to other einops calls to propagate shape information without manual indexing.

```python
import numpy as np
from einops import rearrange, parse_shape

x = np.zeros([2, 3, 5, 7])

# --- Extract named dimensions ---
dims = parse_shape(x, 'batch _ h w')
print(dims)   # {'batch': 2, 'h': 5, 'w': 7}

# --- Propagate parsed shape into a subsequent rearrange ---
y = np.zeros([700])
reshaped = rearrange(y, '(b c h w) -> b c h w', **parse_shape(x, 'b _ h w'))
assert reshaped.shape == (2, 10, 5, 7)   # infers c = 700 / (2*5*7) = 10

# --- Validate a specific expected shape ---
try:
    parse_shape(x, 'batch 3 h w')   # raises if dim 1 != 3
except RuntimeError as e:
    print(e)   # "Length of anonymous axis does not match: batch 3 h w (2, 3, 5, 7)"

# parse_shape with ellipsis skips a variable number of middle dims
parse_shape(x, 'batch ... w')   # {'batch': 2, 'w': 7}
```

---

## `asnumpy` — Framework-Agnostic Tensor-to-NumPy Conversion

`asnumpy` converts a tensor from any supported framework (PyTorch, JAX, TensorFlow, CuPy, etc.) to a `numpy.ndarray` without requiring explicit framework imports at the call site.

```python
import numpy as np
from einops import asnumpy

# NumPy passthrough
arr = np.random.randn(3, 4)
assert type(asnumpy(arr)) is np.ndarray

# PyTorch example (requires torch)
# import torch
# t = torch.randn(3, 4)
# np_arr = asnumpy(t)
# assert isinstance(np_arr, np.ndarray) and np_arr.shape == (3, 4)

# JAX example (requires jax)
# import jax.numpy as jnp
# j = jnp.ones((5, 6))
# assert asnumpy(j).shape == (5, 6)
```

---

## `Rearrange` and `Reduce` Layers — Drop-in `nn.Module` Wrappers

`einops.layers` exposes `Rearrange` and `Reduce` as native neural-network layers for PyTorch, TensorFlow, Flax, and Paddle. They accept the same pattern syntax as the functional API and integrate transparently into `Sequential` / model definitions. PyTorch layers are `torch.compile`-compatible and `torch.jit.script`-compatible.

```python
import torch
import torch.nn as nn
from einops.layers.torch import Rearrange, Reduce

# --- CNN classifier: replace manual flatten with Rearrange ---
model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, padding=1),   # (B, 16, H, W)
    nn.ReLU(),
    nn.Conv2d(16, 32, kernel_size=3, padding=1),  # (B, 32, H, W)
    Reduce('b c (h 2) (w 2) -> b c h w', 'max'),  # 2×2 max-pool
    Rearrange('b c h w -> b (c h w)'),             # flatten, no .view()
    nn.Linear(32 * 14 * 14, 10),
)
x = torch.randn(4, 1, 28, 28)
print(model(x).shape)  # torch.Size([4, 10])

# --- Vision Transformer patch embedding ---
patch_embed = nn.Sequential(
    # reshape image into non-overlapping 8×8 patches
    Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=8, p2=8),
    nn.Linear(8 * 8 * 3, 256),
)
img = torch.randn(2, 3, 64, 64)
print(patch_embed(img).shape)  # torch.Size([2, 64, 256])

# --- Enable torch.compile for functional einops (einops < 0.7 / torch < 2.4) ---
from einops._torch_specific import allow_ops_in_compiled_graph
allow_ops_in_compiled_graph()

compiled_model = torch.compile(model)
print(compiled_model(x).shape)  # torch.Size([4, 10])
```

---

## `EinMix` Layer — Generalized Linear Projection

`EinMix` (available in `einops.layers.torch`, `.tensorflow`, `.flax`, `.paddle`) is a learnable linear layer that mixes axes defined by `weight_shape`. It combines an einsum (projection) with optional pre- and post-`rearrange` steps, enabling projections along any axis without manual transpositions.

```python
import torch
from einops.layers.torch import EinMix

# --- Standard channel-wise linear (equivalent to nn.Linear applied to last dim) ---
linear = EinMix('t b cin -> t b cout', weight_shape='cin cout',
                bias_shape='cout', cin=64, cout=128)
x = torch.randn(10, 4, 64)   # time=10, batch=4, cin=64
print(linear(x).shape)       # torch.Size([10, 4, 128])

# --- Mix along the spatial height axis, not channels ---
height_mix = EinMix('h w c -> hout w c', weight_shape='h hout',
                    bias_shape='hout', h=32, hout=32)
feat = torch.randn(32, 32, 64)
print(height_mix(feat).shape)   # torch.Size([32, 32, 64])

# --- Channel-wise scaling (like a learnable per-channel gain) ---
scale = EinMix('t b c -> t b c', weight_shape='c', c=128)
print(scale(torch.randn(10, 4, 128)).shape)  # torch.Size([10, 4, 128])

# --- Multi-head projection: each head has its own weight ---
mh_proj = EinMix(
    't b (head cin) -> t b (head cout)',
    weight_shape='head cin cout',
    bias_shape='head cout',
    head=8, cin=32, cout=32,
)
tokens = torch.randn(20, 4, 256)  # 8 heads × 32 dim
print(mh_proj(tokens).shape)      # torch.Size([20, 4, 256])
```

---

## Array API Standard Support — Framework-Agnostic Operations

`einops.array_api` exposes `rearrange`, `reduce`, `repeat`, `pack`, and `unpack` for any tensor type that implements the [Python Array API standard](https://data-apis.org/array-api/latest/), including NumPy ≥ 2.0, MLX, CuPy, JAX, and sparse/distributed tensors.

```python
# Works with any Array API-compliant library
import numpy as np
from einops.array_api import rearrange, reduce, pack, unpack

# NumPy 2.0 Array API namespace
xp = np
x = xp.zeros((4, 32, 32, 3))

# Identical patterns to the standard einops API
out = rearrange(x, 'b h w c -> b c h w')
assert out.shape == (4, 3, 32, 32)

pool = reduce(x, 'b (h h2) (w w2) c -> b h w c', 'mean', h2=2, w2=2)
assert pool.shape == (4, 16, 16, 3)

# MLX (Apple Silicon) — same code
# import mlx.core as mx
# x_mlx = mx.zeros((4, 32, 32, 3))
# from einops.array_api import rearrange as rearrange_api
# out_mlx = rearrange_api(x_mlx, 'b h w c -> b c h w')
```

---

## Summary

einops addresses the core pain points of multi-dimensional tensor manipulation: cryptic shape errors, inconsistent API behavior across frameworks, and unreadable reshape/transpose chains. Its pattern language serves as both runtime assertion and documentation, making code immediately understandable. Typical use cases span computer vision (patch embeddings, pooling, space-to-depth), sequence modelling (attention score computation, token packing for multi-modal transformers), and scientific computing (arbitrary-dimensional reductions). The library is particularly valuable in research settings where models are rapidly iterated and tensor shapes change frequently.

Integration is straightforward: install `einops`, import the required functions or layers, and use them anywhere framework-native operations appear. The functional API (`rearrange`, `reduce`, `repeat`, `einsum`, `pack`, `unpack`) is stateless and works with plain tensors from any supported library. The layer API (`Rearrange`, `Reduce`, `EinMix`) slots directly into `nn.Sequential` or model `__init__` blocks and supports `torch.compile`, making it zero-overhead in production PyTorch workflows. For projects targeting multiple backends (e.g., NumPy and JAX), `einops.array_api` provides the same interface without any backend-specific code.