arogozhnikov/einops

★ 9,485PythonAudience · researcherComplexity · 2/5Setup · easy

Mindmap

mindmap
  root((einops))
    What it does
      Tensor reshaping
      Dimension rearranging
      Readable patterns
    Core Functions
      rearrange
      reduce
      repeat
    Frameworks
      PyTorch
      NumPy
      JAX
      TensorFlow
    Use Cases
      ML model code
      Research scripts
      Framework-agnostic ops

mindmap root((einops)) What it does Tensor reshaping Dimension rearranging Readable patterns Core Functions rearrange reduce repeat Frameworks PyTorch NumPy JAX TensorFlow Use Cases ML model code Research scripts Framework-agnostic ops

Click or tap to explore — scroll the page freely

Things people build with this

USE CASE 1

Rewrite confusing NumPy reshape calls in your ML code as readable einops patterns that explain the transformation.

USE CASE 2

Use the same tensor rearrangement code across a PyTorch model and a JAX experiment without rewriting it.

USE CASE 3

Add einops layers directly inside a neural network definition to make dimension transformations explicit and self-documenting.

Tech stack

PythonNumPyPyTorchJAXTensorFlowPaddleMLX

Getting it running

Difficulty · easy Time to first run · 5min

In plain English

Einops is a Python library that makes it easier to write and read code that reshapes, transposes, or combines multi-dimensional arrays (called tensors) used in machine learning. In machine learning, data is constantly being sliced, rearranged, and combined across many dimensions, and the standard ways to do this tend to produce code that is hard to follow. Einops offers a text-based notation where you write out the shape transformation you want in a readable pattern, and the library figures out how to perform it. The core idea borrows from a mathematical notation called Einstein summation, where you describe an operation by naming the dimensions explicitly rather than passing a series of numbers and flags. For example, instead of writing a reshape command with magic numbers, you write something like 'batch channels height width -> batch (channels height width)' and the meaning is clear from the pattern itself. The three main functions are rearrange (for reordering and reshaping), reduce (for reducing dimensions by averaging or summing), and repeat (for expanding a tensor along a new axis). The library works with the major Python array frameworks: NumPy, PyTorch, JAX, TensorFlow, Paddle, MLX, tinygrad, and others. This means the same einops notation works regardless of which underlying framework a project uses. It also provides layer objects so you can slot these operations directly into a neural network model definition. Installation is a single pip command. The project ships four tutorial notebooks covering basics, deep learning use cases, packing and unpacking multiple arrays, and PyTorch-specific examples. A companion website hosts documentation and testimonials, and the project has been adopted in over 10,000 other GitHub repositories. A related paper was accepted at ICLR 2022 and received an oral presentation, which is a competitive distinction in the machine learning research community. The project is actively maintained with recent updates adding new framework backends and compatibility with PyTorch's compilation tools.

Copy-paste prompts

Prompt 1

Using einops rearrange, convert a batch of images with shape (batch, height, width, channels) to (batch, channels, height, width) for a PyTorch model.

Prompt 2

Show me how to use einops reduce to average-pool a feature map from (batch, channels, h, w) to (batch, channels) in a neural network.

Prompt 3

Rewrite this NumPy code using einops so it is more readable: x.transpose(0,2,1).reshape(batch, -1).

Prompt 4

How do I use einops repeat to expand a single image tensor into a batch of 8 identical copies?

Open on GitHub → Explain another repo

← arogozhnikov on gitmyhub — every repo by this author, as a profile.

Verify against the repo before relying on details.