DOCSEARCHbyAlgolia
Submit Docs

einops Documentation

repository·main·Indexed 27 days ago

https://github.com/arogozhnikov/einops

A library for flexible and powerful tensor operations using Einstein-inspired notation. It provides framework-independent manipulation tools—including rearrange, reduce, repeat, pack, and unpack—that work across numpy, pytorch, jax, mlx, tensorflow, and other libraries implementing the Python array API standard. The library allows for readable tensor reshaping, transposition, and reduction, and includes specialized layers for neural network architectures.

Tokens
12.7K
Snippets
30
Records
77
Agent score
93%

What's inside einops

  1. Supported frameworks for einops

    main

    Einops is compatible with a wide range of tensor libraries. It supports frameworks that implement the Python array API standard, including:

    • Core support: numpy, pytorch, tensorflow, jax, cupy.
    • Community support: flax, paddle, oneflow, tinygrad, pytensor.
    • Other compatible libraries: MLX (Apple), pydata/sparse (>= 0.15), cubed (distributed tensors), quantco/ndonnx, and dask (via array-api-compat).

    When using the array API standard, it is recommended to import from einops.array_api instead of the top-level einops module.

    # Instead of:
from einops import rearrange

# Use:
from einops.array_api import rearrange

  5. Explore einops implementations in real-world projects

    main

    To see how einops is applied in practice across different domains, you can examine the following implementations:

  6. Run einops tests for specific frameworks

    main

    Starting from version 0.8.1, einops distributes its tests as part of the package. You can run tests against specific frameworks (e.g., numpy, pytorch, jax). Every framework is tested against numpy, which is a requirement for running tests.

    Use the python -m einops.tests.run_tests command. If you want to install the required dependencies for the chosen frameworks into your current virtual environment, include the --pip-install flag.

    # pip install einops pytest
python -m einops.tests.run_tests numpy pytorch jax --pip-install

  7. Access Einops tutorials via Jupyter notebooks

    main

    You can learn how to use einops through a series of structured notebooks. These are available directly on GitHub or via nbviewer for easier viewing.

    Tutorial Series:

    • Part 1: Einops Basics - Fundamental concepts and syntax.
    • Part 2: Einops for Deep Learning - Applying einops to neural network architectures.
    • Part 3: EinMix for great MLPs - Using EinMix in Multi-Layer Perceptrons.
    • Part 4: einops.pack and einops.unpack - Advanced tensor manipulation using packing and unpacking operations.
    https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb
https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb
https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb
https://nbviewer.jupyter.org/github/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

  11. Enable einops operations in torch.compile graphs

    main

    To ensure einops operations like rearrange, reduce, repeat, einsum, pack, and unpack are compatible with torch.compile (TorchDynamo), the library automatically attempts to register them using torch._dynamo.allow_in_graph.

    This registration happens automatically upon importing the module. If you are using PyTorch 2.0 or higher, these operations should be recognized as graph-compatible. For PyTorch versions 2.8 and above, explicit registration is no longer required by einops.

  12. Use einops layers in PyTorch nn.Sequential

    main

    You can use Rearrange and Reduce from einops.layers.torch directly within nn.Sequential blocks to replace manual .view(), .reshape(), or .transpose() calls. This makes the model architecture more readable and provides explicit shape validation, raising errors if input dimensions do not match the expected pattern.

    Example of replacing a manual view in a ConvNet:

    from einops.layers.torch import Rearrange
import torch.nn as nn

conv_net_new = nn.Sequential(
    nn.Conv2d(1, 10, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Conv2d(10, 20, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Dropout2d(),
    Rearrange('b c h w -> b (c h w)'),
    nn.Linear(320, 50),
    nn.ReLU(),
    nn.Dropout(),
    nn.Linear(50, 10),
    nn.LogSoftmax(dim=1)
)
    from einops.layers.torch import Rearrange
import torch.nn as nn

conv_net_new = nn.Sequential(
    nn.Conv2d(1, 10, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Conv2d(10, 20, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Dropout2d(),
    Rearrange('b c h w -> b (c h w)'),
    nn.Linear(320, 50),
    nn.ReLU(),
    nn.Dropout(),
    nn.Linear(50, 10),
    nn.LogSoftmax(dim=1)
)