Einops: A New Flavor of Deep Learning Operations

0.8.2 · active · verified Sun Mar 29

Einops (Einstein operations) is a Python library that provides a flexible and powerful way to reshape and manipulate tensors in deep learning frameworks like PyTorch, TensorFlow, JAX, and NumPy. It simplifies complex tensor operations using a human-readable notation, often replacing verbose permutations, transpositions, and reshaping operations. The library is actively maintained with frequent releases, currently at version 0.8.2.

Warnings

breaking TensorFlow layers in `einops` were updated in v0.8.0 to align with TF 2.16+ and are no longer compatible with older TensorFlow versions (e.g., TF 2.13).
Fix: Upgrade TensorFlow to 2.16 or newer, or stick to `einops < 0.8.0` for older TensorFlow versions.
breaking Support for Python 3.7 was officially dropped in `einops` v0.7.0. The minimum required Python version is now 3.9.
Fix: Upgrade your Python environment to 3.9 or newer.
breaking As of v0.8.2, the minimum required Python version for `einops` is Python 3.9.
Fix: Ensure your Python environment is 3.9 or newer before upgrading to `einops >= 0.8.2`.
deprecated Support for the Gluon (MXNet) backend was dropped in v0.6.1 and confirmed removed in v0.7.0. Operations with MXNet tensors are no longer supported directly by `einops`.
Fix: Migrate to a supported backend like PyTorch, TensorFlow, JAX, or NumPy.
gotcha The integration with `torch.compile` has evolved across versions. In `einops < 0.7.0`, explicit registration via `einops._torch_specific.allow_ops_in_compiled_graph()` was required. From `v0.7.0` onwards, `torch.compile` integration became largely automatic. With `einops >= 0.8.2` and `torch >= 2.8`, `torch.compile` natively handles `einops` operations without any specific `einops` hints or registrations.
Fix: For `einops < 0.7.0`, ensure `einops._torch_specific.allow_ops_in_compiled_graph()` is called. For `einops >= 0.7.0`, no explicit action is usually needed, but ensure your PyTorch version is recent enough (especially `torch >= 2.8` for `einops >= 0.8.2`) to benefit from the latest native compilation.
gotcha Ellipsis (`...`) support was added to `EinMix` layers in v0.8.1, allowing for more flexible input patterns in mixed-precision operations.
Fix: Upgrade to `einops >= 0.8.1` to utilize ellipsis in `EinMix` patterns.

Install

pip install einops Install stable version

Imports

rearrange
```
from einops import rearrange
```
reduce
```
from einops import reduce
```
repeat
```
from einops import repeat
```
einsum
```
from einops import einsum
```
pack
```
from einops import pack
```
unpack
```
from einops import unpack
```
EinMix
```
from einops.layers.torch import EinMix
```
Import path for layers depends on the specific deep learning framework (e.g., `einops.layers.tf`, `einops.layers.jax`).

Quickstart

This quickstart demonstrates the core `rearrange` operation using NumPy. It shows how to combine dimensions (e.g., batch and height) and flatten dimensions (e.g., width and channel) using a concise, readable pattern string.

import numpy as np
from einops import rearrange

# Suppose we have a batch of 6 images, each 96x96 with 3 color channels
images = np.random.randn(6, 96, 96, 3)
print(f"Original shape: {images.shape}")

# Rearrange to stack images vertically (batch and height become one dimension)
stacked_images = rearrange(images, 'b h w c -> (b h) w c')
print(f"Stacked shape: {stacked_images.shape}")

# Alternatively, flatten width and channel for a 2D representation
flattened_data = rearrange(images, 'b h w c -> b h (w c)')
print(f"Flattened shape: {flattened_data.shape}")

view raw JSON →