einops

https://github.com/arogozhnikov/einops
Admin
einops provides flexible and powerful tensor operations for readable and reliable code, with support...
Tokens:73,658
Snippets:541
Trust Score:9.2
License:MIT
Update:4 days ago
Show doc for...
Tokens:
Raw
### Install einops, numpy, pillow

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Installs required libraries for the tutorial using pip.

```python
# we need some libraries for this demo
%pip install einops numpy pillow -q
```

--------------------------------

### Install numpy and einops

Source: https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

Install the required libraries for the examples. The output notes that the kernel may need to be restarted to use updated packages.

```python
# install necessary libraries
%pip install numpy einops -q
```

--------------------------------

### Run einops tests with numpy, pytorch, jax

Source: https://github.com/arogozhnikov/einops/blob/main/README.md

Executes the test suite for specified frameworks. The frameworks list is an example; any subset can be provided. Use --pip-install to install requirements in the current virtualenv; omit it if dependencies are already installed.

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

--------------------------------

### Install einops with pip

Source: https://github.com/arogozhnikov/einops/blob/main/README.md

Install the einops library using pip. The command is plain and simple; uv pip install einops also works.

```bash
pip install einops
```

--------------------------------

### Setup numpy display as images

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Imports numpy and enables image display for numpy arrays in Jupyter notebooks.

```python
# Examples are given for numpy. This code also setups ipython/jupyter/jupyterlite
# so that numpy arrays in the output are displayed as images
import numpy
from utils import display_np_arrays_as_images

display_np_arrays_as_images()
```

--------------------------------

### Output of backpropagation

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows the output of the backpropagation example: tensor(320., dtype=torch.float64).

```text
Output:
tensor(320., dtype=torch.float64)

```

--------------------------------

### Import ResNet blocks

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Imports BasicBlock, Bottleneck, and ResNet from torchvision.models.resnet for use in the examples.

```python
from torchvision.models.resnet import BasicBlock, Bottleneck, ResNet
```

--------------------------------

### Setup backend and convert x

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Prints the selected backend and initializes the tensor for the chosen framework. For tensorflow, it creates a gradient tape and converts x to a variable; for pytorch, it converts x to a torch tensor with requires_grad=True.

```python
print(f"selected {flavour} backend")
if flavour == "tensorflow":
    import tensorflow as tf

    tape = tf.GradientTape(persistent=True)
    tape.__enter__()
    x = tf.Variable(x) + 0
else:
    assert flavour == "pytorch"
    import torch

    x = torch.from_numpy(x)
    x.requires_grad = True
```

--------------------------------

### Usage example for super-resolution (commented)

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Commented-out code showing how the original super-resolution example was used with PIL and torchvision. It loads an image, converts to YCbCr, applies the model, and merges channels back. Note the comment about remembering the order of components in PIL.Image.size.

```python
## that's how this code was mentioned to use

# from PIL import Image
# img = Image.open(opt.input_image).convert('YCbCr')
# y, cb, cr = img.split()

# model = torch.load(opt.model)
# img_to_tensor = ToTensor()
# input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])

# if opt.cuda:
#     model = model.cuda()
#     input = input.cuda()

# out = model(input)
# out = out.cpu()
# out_img_y = out[0].detach().numpy()
# out_img_y *= 255.0
# out_img_y = out_img_y.clip(0, 255)
# out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')

# out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)
# out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
# out_img = Image.merge('YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')

## Benefits

# - no need to remembder the order of components in PIL.Image.size (as you see, it is actually different)
# - code explicitly shows shapes passed in and out
# - normalization to [0, 1] range and back is also explicit (it is needed to rememebed in original code that division by 255 is done by ToTensor)
```

--------------------------------

### Build a PyTorch model with Reduce layer

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Example of building a Sequential model in PyTorch using the Reduce layer for combined pooling and flattening, with a comment about Flax syntax.

```python
# example given for pytorch, but code in other frameworks is almost identical
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Reduce

model = Sequential(
    Conv2d(3, 6, kernel_size=5),
    MaxPool2d(kernel_size=2),
    Conv2d(6, 16, kernel_size=5),
    # combined pooling and flattening in a single step
    Reduce('b c (h 2) (w 2) -> b (c h w)', 'max'), 
    Linear(16*5*5, 120), 
    ReLU(),
    Linear(120, 10), 
    # In flax, the {'axis': value} syntax for specifying values for axes is mandatory:
    # Rearrange('(b1 b2) d -> b1 b2 d', {'b1': 12}), 
)
```

--------------------------------

### Preparations for ViT example

Source: https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

Creates random patch tokens and zero class tokens for the vision transformer example. The shapes are batch=6, height=16, width=16, channels=256.

```python
# preparations
batch, height, width, c = 6, 16, 16, 256
patch_tokens = np.random.random([batch, height, width, c])
class_tokens = np.zeros([batch, c])
```

--------------------------------

### Reduce: Compute Result Shape

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Another reduce example. The result shape is not given; the reader is asked to compute it.

```python
# yet another example. Can you compute result shape?
reduce(ims, "(b1 b2) h w c -> (b2 h) (b1 w)", "mean", b1=2)
```

--------------------------------

### Examples of einops.rearrange usage

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/rearrange.md

Shows various rearrangement patterns: stacking, transposing, concatenating, flattening, splitting, and space-to-depth. Each example uses a list of 32 images with shape (30, 40, 3) and outputs the resulting shape.

```python
# suppose we have a set of 32 images in "h w c" format (height-width-channel)
>>> images = [np.random.randn(30, 40, 3) for _ in range(32)]

# stack along first (batch) axis, output is a single array
>>> rearrange(images, 'b h w c -> b h w c').shape
(32, 30, 40, 3)

# stacked and reordered axes to "b c h w" format
>>> rearrange(images, 'b h w c -> b c h w').shape
(32, 3, 30, 40)

# concatenate images along height (vertical axis), 960 = 32 * 30
>>> rearrange(images, 'b h w c -> (b h) w c').shape
(960, 40, 3)

# concatenated images along horizontal axis, 1280 = 32 * 40
>>> rearrange(images, 'b h w c -> h (b w) c').shape
(30, 1280, 3)

# flattened each image into a vector, 3600 = 30 * 40 * 3
>>> rearrange(images, 'b h w c -> b (c h w)').shape
(32, 3600)

# split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
>>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
(128, 15, 20, 3)

# space-to-depth operation
>>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
(32, 15, 20, 12)
```

--------------------------------

### Basic max-reduction examples

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/reduce.md

Examples showing max-reduction over the first axis using both short and verbose axis names. The axis not appearing on the RHS is reduced over.

```python
>>> x = np.random.randn(100, 32, 64)

# perform max-reduction on the first axis
# Axis t does not appear on RHS - thus we reduced over t
>>> y = reduce(x, 't b c -> b c', 'max')

# same as previous, but using verbose names for axes
>>> y = reduce(x, 'time batch channel -> batch channel', 'max')
```

--------------------------------

### Full usage with both models

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Complete example loading an image, applying both models, and converting outputs. Uses torch.from_numpy and rearrange for the new model. Note the TODO comment about numpy.asarray.

```python
input_image = '../../logo/einops_logo_350x350.png'
from PIL import Image
import numpy as np

from torchvision.transforms import ToTensor
model = SuperResolutionNetOld(upscale_factor=2)

img = Image.open(input_image).convert('YCbCr')
y, cb, cr = img.split()

img_to_tensor = ToTensor()
input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])
out = model(input)

out_img_y = out[0].detach().numpy()
out_img_y = np.clip(out_img_y[0] * 255, 0, 255)

model = SuperResolutionNetNew(upscale_factor=2)

img = Image.open(input_image).convert('YCbCr')
y, cb, cr = img.split()

# TODO numpy.asarray
y = torch.from_numpy(np.array(y, dtype='float32') / 255)
out = model(rearrange(y, 'h w -> () () h w'))

out_img_y = asnumpy(rearrange(out, '() h w -> h w'))
out_img_y = np.clip(out_img_y * 255, 0, 255)
```

--------------------------------

### Squeeze and unsqueeze (expand_dims) with rearrange

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows how to add and remove singleton axes using `()` in the `rearrange` pattern. The example extracts an image from a batch, adds a dummy batch axis, flattens spatial and channel axes, then removes the dummy axis to get class predictions.

```python
# models typically work only with batches,
# so to predict a single image ...
image = rearrange(x[0, :3], "c h w -> h w c")
# ... create a dummy 1-element axis ...
y = rearrange(image, "h w c -> () c h w")
# ... imagine you predicted this with a convolutional network for classification,
# we'll just flatten axes ...
predictions = rearrange(y, "b c h w -> b (c h w)")
# ... finally, decompose (remove) dummy axis
predictions = rearrange(predictions, "() classes -> classes")
```

--------------------------------

### Rearrange with verbose axis names

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Same rearrangement as the previous example but using verbose axis names (height, width, color).

```python
# we could use more verbose names for axes, and result is the same:
rearrange(ims[0], "height width color -> width height color")
# when you operate on same set of axes many times,
# you usually come up with short names.
# That's what we do throughout tutorial - we'll use b (for batch), h, w, and c
```

--------------------------------

### Benchmark setup for attention implementations

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Creates input tensors K, V, Q on CUDA and times both the class-based and einsum-based attention implementations using %timeit. The timing includes torch.cuda.synchronize() to ensure accurate GPU timing.

```python
args = dict(
    K=torch.zeros(32, 128, 40).cuda(),
    V=torch.zeros(32, 128, 40).cuda(),
    Q=torch.zeros(32, 128, 30).cuda(),
)

%timeit -n100 result_old = Attention()(**args); torch.cuda.synchronize()
%timeit -n100 result_new = attention(**args); torch.cuda.synchronize()
```

--------------------------------

### Stacking: stack along last dimension

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Alternative stacking example that places the new axis (enumerating tensors) at the last position in the output pattern.

```python
# or maybe stack along last dimension?
tensors = rearrange(list_of_tensors, "b c h w -> h w c b")
guess(tensors.shape)
```

--------------------------------

### Setup and comparison of attention layers

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Initializes two attention layers (old and new) with identical parameters and creates random inputs to compare their outputs. The code sets up the arguments and runs both layers.

```python
n_heads = 8
d_k = 32
d_v = 64
d_model = 100
t = 51
l = 53
batch = 30

layer1 = initialize(MultiHeadAttentionOld(n_heads, d_k=d_k, d_v=d_v, d_model=d_model)).eval().cuda()
layer2 = initialize(MultiHeadAttentionNew(n_heads, d_k=d_k, d_v=d_v, d_model=d_model)).eval().cuda()

args = dict(
    q=torch.randn(batch, l, d_model),
    k=torch.randn(batch, t, d_model) * 0.1,
    v=torch.randn(batch, t, d_model),
    mask=torch.randn(batch, l, t) > 0,
)
args = {k:v.cuda() for k, v in args.items()}

o1, a1 = layer1(**args)
o2, a2 = layer2(**args)

```

--------------------------------

### EinMix patcher and positional embedding (commented)

Source: https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

A commented example showing how to use EinMix to split images into patches and apply positional embedding in a single shot. This is provided as a reference and is not executed.

```python
# below it is assumed that you already
# 1) split batch of images into patches 2) applied linear projection and 3) used positional embedding.

# We'll skip that here. But hey, here is an einops-style way of doing all of that in a single shot!
# from einops.layers.torch import EinMix
# patcher_and_posembedder = EinMix('b (h h2) (w w2) c -> b h w c_out', weight_shape='h2 w2 c c_out',
#                                  bias_shape='h w c_out', h2=..., w2=...)
# patch_tokens_bhwc = patcher_and_posembedder(images_bhwc)
```

--------------------------------

### Initialize both models

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Initializes the old and new super-resolution models with upscale_factor=3. The variable 'initialize' is not defined in the snippet; it is assumed to be a helper function from the original example.

```python
model1 = initialize(SuperResolutionNetOld(upscale_factor=3))
model2 = initialize(SuperResolutionNetNew(upscale_factor=3))
```

--------------------------------

### Adaptive pooling and global average pooling

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/reduce.md

Examples of adaptive 2D max-pooling to a 3x4 grid and global average pooling, showing the resulting tensor shapes.

```python
# adaptive 2d max-pooling to 3 * 4 grid,
# each element is max of 10x10 tile in the original tensor.
>>> reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h1=3, w1=4).shape
(10, 20, 3, 4)

# Global average pooling
>>> reduce(x, 'b c h w -> b c', 'mean').shape
(10, 20)
```

--------------------------------

### Subtracting mean with different reduction patterns

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/reduce.md

Examples of subtracting mean over batch for each channel, per-image mean for each channel, and the equivalent using empty compositions.

```python
# subtracting mean over batch for each channel;
# similar to x - np.mean(x, axis=(0, 2, 3), keepdims=True)
>>> y = x - reduce(x, 'b c h w -> 1 c 1 1', 'mean')

# Subtracting per-image mean for each channel
>>> y = x - reduce(x, 'b c h w -> b c 1 1', 'mean')

# same as previous, but using empty compositions
>>> y = x - reduce(x, 'b c h w -> b c () ()', 'mean')
```

--------------------------------

### Stacking: create list of tensors

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Creates a list of tensors from a batch tensor `x` by iterating over the first axis. This list is used in subsequent stacking examples.

```python
list_of_tensors = list(x)
```

--------------------------------

### Repeat each element (pixel) 3 times by changing parenthesis order

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

By placing 'repeat' inside the parentheses after 'w', each element is repeated consecutively, unlike the previous example. The output is not shown.

```python
# order of axes matters as usual - you can repeat each element (pixel) 3 times
# by changing order in parenthesis
repeat(ims[0], "h w c -> h (w repeat) c", repeat=3)
```

--------------------------------

### Import einops and PyTorch modules

Source: https://github.com/arogozhnikov/einops/blob/main/docs/pytorch-examples.html

Imports PyTorch, NumPy, math, and einops functions (rearrange, reduce, asnumpy, parse_shape) and layers (Rearrange, Reduce). Required for all examples on this page.

```python
# start from importing some stuff
import torch
import torch.nn as nn 
import torch.nn.functional as F
import numpy as np
import math

from einops import rearrange, reduce, asnumpy, parse_shape
from einops.layers.torch import Rearrange, Reduce
```

--------------------------------

### Rearrange with numpy array

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Applies the rearrange operation to a numpy array 'ims' with the pattern 'b h w c -> w (b h) c'. This is the base example before adapting to other frameworks.

```python
rearrange(ims, 'b h w c -> w (b h) c')
```

--------------------------------

### Serve documentation with hatch

Source: https://github.com/arogozhnikov/einops/blob/main/README.md

Builds and serves the documentation locally at http://localhost:8000/.

```bash
hatch run docs:serve  # Serving on http://localhost:8000/
```

--------------------------------

### Display first image

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Displays the first image from the loaded batch.

```python
# display the first image (whole 4d tensor can't be rendered)
ims[0]
```

--------------------------------

### Instantiate and apply Rearrange/Reduce layers

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows how to create Rearrange and Reduce layers with the same parameters as operations, and apply them to a tensor.

```python
layer = Rearrange(pattern, **axes_lengths)
layer = Reduce(pattern, reduction, **axes_lengths)

# apply layer to tensor
x = layer(x)
```

--------------------------------

### Output of backend selection

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows the output when the flavour is set to 'pytorch': it prints 'selected pytorch backend'.

```text
Output:
selected pytorch backend

```

--------------------------------

### Formula for einsum example

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/einsum.md

Mathematical formula for the einsum example: output[a,b,k] = sum over c,d,g of x[a,b,c]*y[c,b,d]*z[a,g,k].

```tex
output[a, b, k] = \sum_{c, d, g} x[a, b, c] * y[c, b, d] * z[a, g, k]
```

--------------------------------

### einops.pack example with numpy

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/pack_unpack.md

Example usage of einops.pack with numpy arrays. It packs tensors of different shapes into one, returning the packed tensor and packed shapes (PS). The axes i, j, k are matched by order, and the rest are concatenated.

```python
>>> from numpy import zeros as Z
>>> inputs = [Z([2, 3, 5]), Z([2, 3, 7, 5]), Z([2, 3, 7, 9, 5])]
>>> packed, ps = pack(inputs, 'i j * k')
>>> packed.shape, ps
((2, 3, 71, 5), [(), (7,), (7, 9)])
```

--------------------------------

### Select framework flavour

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Sets the 'flavour' variable to either 'tensorflow' or 'pytorch' to choose the backend for the tutorial.

```python
# select "tensorflow" or "pytorch"
flavour = "pytorch"
```

--------------------------------

### Display second image

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Displays the second image from the loaded batch.

```python
# second image in a batch
ims[1]
```

--------------------------------

### Rearrange adapted for multiple frameworks

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Shows the same rearrange operation adapted for pytorch, tensorflow, gluon, cupy, jax, and paddle. The code is identical across frameworks, demonstrating einops' universal notation.

```python
# pytorch:
rearrange(ims, 'b h w c -> w (b h) c')
# tensorflow:
rearrange(ims, 'b h w c -> w (b h) c')
# gluon:
rearrange(ims, 'b h w c -> w (b h) c')
# cupy:
rearrange(ims, 'b h w c -> w (b h) c')
# jax:
rearrange(ims, 'b h w c -> w (b h) c')
# paddle:
rearrange(ims, 'b h w c -> w (b h) c')
```

--------------------------------

### Disproportionate resize

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Reduces with mean to perform a disproportionate resize.

```python
reduce(ims, "b (h 4) (w 3) c -> (h) (b w)", "mean")
```

--------------------------------

### Create fake batch and import matplotlib and torchvision utils

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Sets up a random fake batch of 100 images with shape [100, 3, 32, 32] and imports matplotlib's pyplot and torchvision.utils for grid visualization.

```python
fake_batch = torch.rand([100, 3, 1, 1]) + torch.zeros([100, 3, 32, 32])
from matplotlib import pyplot as plt
import torchvision.utils as vutils
```

--------------------------------

### parse_shape usage example

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/parse_shape.md

Shows how to use parse_shape to extract axis lengths from a tensor, skipping dimensions with underscore. The output can be used as axes_lengths for other operations like rearrange.

```python
# Use underscore to skip the dimension in parsing.
>>> x = np.zeros([2, 3, 5, 7])
>>> parse_shape(x, 'batch _ h w')
{'batch': 2, 'h': 5, 'w': 7}

# `parse_shape` output can be used to specify axes_lengths for other operations:
>>> y = np.zeros([700])
>>> rearrange(y, '(b c h w) -> b c h w', **parse_shape(x, 'b _ h w')).shape
(2, 10, 5, 7)
```

--------------------------------

### TODO comments for future work

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Placeholder comments indicating planned additions for hierarchical softmax and reinforcement learning content.

```python
# TODO: Hierarchical softmax
# TODO: some reinforcement stuff would also be needed
```

--------------------------------

### 2D max-pooling with kernel size

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/reduce.md

Examples of 2D max-pooling with kernel size 2x2 for image processing, using both named and anonymous axes. Anonymous axes can only be used for reduced axes.

```python
>>> x = np.random.randn(10, 20, 30, 40)

# 2d max-pooling with kernel size = 2 * 2 for image processing
>>> y1 = reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h2=2, w2=2)

# same as previous, using anonymous axes,
# note: only reduced axes can be anonymous
>>> y1 = reduce(x, 'b c (h1 2) (w1 2) -> b c h1 w1', 'max')
```

--------------------------------

### Import Rearrange and Reduce layers for multiple backends

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows how to import the Rearrange and Reduce layers from einops for PyTorch, Flax, TensorFlow, and Chainer backends.

```python
from einops.layers.torch import Rearrange, Reduce
from einops.layers.flax import Rearrange, Reduce
from einops.layers.tensorflow import Rearrange, Reduce
from einops.layers.chainer import Rearrange, Reduce
```

--------------------------------

### Using Layers in a PyTorch Model

Source: https://github.com/arogozhnikov/einops/blob/main/README.md

Use Rearrange layer to flatten tensors within a Sequential model without writing a forward method. This example is for PyTorch, but code in other frameworks is almost identical.

```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),
    # flattening without need to write forward
    Rearrange('b c h w -> b (c h w)'),
    Linear(16*5*5, 120),
    ReLU(),
    Linear(120, 10),
)
```

--------------------------------

### Load test images

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Loads a batch of images from a .npy file and prints its shape and dtype.

```python
ims = numpy.load("./resources/test_images.npy", allow_pickle=False)
# There are 6 images of shape 96x96 with 3 color channels packed into tensor
print(ims.shape, ims.dtype)
```

--------------------------------

### EinMix counter-part for linear layer

Source: https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb

EinMix layer that performs the same operation as the einsum example, but handles parameter initialization and management, and includes a bias term.

```python
mix_channels = Mix('t b c -> t b c_out', weight_shape='c c_out', bias_shape='c_out', ...)
result = mix_channels(embeddings)
```

--------------------------------

### Instantiate ShuffleUnitOld and ShuffleUnitNew

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Creates two ShuffleUnit models with identical parameters to compare the old and new implementations.

```python
model1 = ShuffleUnitOld(32, 32, groups=4, grouped_conv=True, combine='add')
model2 = ShuffleUnitNew(32, 32, groups=4, grouped_conv=True, combine='add')
```

--------------------------------

### Benchmark output for attention implementations

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Shows the timing results from the %timeit commands. Both implementations run at approximately 336 µs per loop, indicating similar performance.

```text
Output:
336 µs ± 7.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
336 µs ± 7.48 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

--------------------------------

### Import numpy for demo

Source: https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

Import numpy for demonstration purposes. The operations work the same way with other frameworks.

```python
# we'll use numpy for demo purposes
# operations work the same way with other frameworks
import numpy as np
```

--------------------------------

### Import einops operations

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Imports the three core einops operations: rearrange, reduce, and repeat.

```python
# we'll use three operations
from einops import rearrange, reduce, repeat
```

--------------------------------

### Rearrange: Order of Axes in Composition

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Shows how the order of axes in composition affects the result. The leftmost axis is most significant, similar to digits in a number. Compare with the next example to see the difference.

```python
# compare with the next example
rearrange(ims, "b h w c -> h (b w) c")
```

```python
# order of axes in composition is different
# rule is just as for digits in the number: leftmost digit is the most significant,
# while neighboring numbers differ in the rightmost axis.

# you can also think of this as lexicographic sort
rearrange(ims, "b h w c -> h (w b) c")
```

--------------------------------

### Repeat operation examples with numpy

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/repeat.md

Shows various repeat patterns: adding a channel dimension, repeating along height, repeating along both axes, upsampling, and pixelation. Uses C-order enumeration when composing axes.

```python
# a grayscale image (of shape height x width)
>>> image = np.random.randn(30, 40)

# change it to RGB format by repeating in each channel
>>> repeat(image, 'h w -> h w c', c=3).shape
(30, 40, 3)

# repeat image 2 times along height (vertical axis)
>>> repeat(image, 'h w -> (repeat h) w', repeat=2).shape
(60, 40)

# repeat image 2 time along height and 3 times along width
>>> repeat(image, 'h w -> (h2 h) (w3 w)', h2=2, w3=3).shape
(60, 120)

# convert each pixel to a small square 2x2, i.e. upsample an image by 2x
>>> repeat(image, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape
(60, 80)

# 'pixelate' an image first by downsampling by 2x, then upsampling
>>> downsampled = reduce(image, '(h h2) (w w2) -> h w', 'mean', h2=2, w2=2)
>>> repeat(downsampled, 'h w -> (h h2) (w w2)', h2=2, w2=2).shape
(30, 40)
```

--------------------------------

### Define a placeholder 2D convolution function

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Defines a simple 2D convolution function that returns the input unchanged, used as a stand-in for actual convolution layers in later examples.

```python
def convolve_2d(x):
    # imagine we have a simple 2d convolution with padding,
    # so output has same shape as input.
    # Sorry for laziness, use imagination!
    return x
```

--------------------------------

### Split in small patches and transpose each patch

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Rearranges images by splitting into small patches and transposing each patch, using h2=8 and w2=8.

```python
rearrange(ims, "b (h1 h2) (w1 w2) c -> (h1 w2) (b w1 h2) c", h2=8, w2=8)
```

--------------------------------

### Split each image in two halves and compute mean

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Reduces with mean after splitting each image into two halves, using h1=2.

```python
reduce(ims, "b (h1 h2) w c -> h2 (b w)", "mean", h1=2)
```

--------------------------------

### Basic einsum example with three tensors

Source: https://github.com/arogozhnikov/einops/blob/main/docs_src/api/einsum.md

Computes a tensor product with three tensors using the pattern 'a b c, c b d, a g k -> a b k'. Summation over axes not on the right-hand side (c, d, g) is performed.

```python
>>> x, y, z = np.random.randn(3, 20, 20, 20)
>>> output = einsum(x, y, z, "a b c, c b d, a g k -> a b k")
```

--------------------------------

### Rearrange to swap axes

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Rearranges images to swap height and width axes.

```python
rearrange(ims, "b h w c -> w (b h) c")
```

--------------------------------

### Transpose with numpy vs einops rearrange

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Compares a raw numpy transpose with the equivalent einops rearrange call, showing how einops makes axis reordering more readable.

```python
y = x.transpose(0, 2, 3, 1)
```

```python
y = rearrange(x, 'b c h w -> b h w c')
```

--------------------------------

### Compare old and new ResNet outputs

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Creates both models with BasicBlock and layers=[2,2,2,3], initializes them, and asserts their outputs are close within atol=1e-3.

```python
x = torch.randn(2, 3, 224, 224)
with torch.no_grad():
    model_old = ResNetOld(BasicBlock, layers=[2, 2, 2, 3])
    model_new = ResNetNew(BasicBlock, layers=[2, 2, 2, 3])
    initialize(model_old)
    initialize(model_new)
    assert torch.allclose(model_old(x), model_new(x), atol=1e-3)
```

--------------------------------

### Simple Attention: Class-based implementation (left)

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Defines an Attention module using torch.bmm and F.softmax. The forward method takes K, V, Q tensors and returns the concatenated result along dim=1. This is the 'left' example in the notebook, used for comparison with the einsum version.

```python
#left
class Attention(nn.Module):
    def __init__(self):
        super(Attention, self).__init__()

    def forward(self, K, V, Q):
        A = torch.bmm(K.transpose(1,2), Q) / np.sqrt(Q.shape[1])
        A = F.softmax(A, 1)
        R = torch.bmm(V, A)
        return torch.cat((R, Q), dim=1)
```

--------------------------------

### Transformer attention: ScaledDotProductAttention and MultiHeadAttentionOld (left)

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Defines ScaledDotProductAttention and MultiHeadAttentionOld modules for multi-head attention. The ScaledDotProductAttention computes scaled dot-product attention with optional masking and dropout. MultiHeadAttentionOld projects inputs, splits into heads, applies attention, and combines with residual connection and layer norm. This is the 'left' example for comparison with the einsum version.

```python
#left
class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)
        self.softmax = nn.Softmax(dim=2)

    def forward(self, q, k, v, mask=None):

        attn = torch.bmm(q, k.transpose(1, 2))
        attn = attn / self.temperature

        if mask is not None:
            attn = attn.masked_fill(mask, -np.inf)

        attn = self.softmax(attn)
        attn = self.dropout(attn)
        output = torch.bmm(attn, v)

        return output, attn



class MultiHeadAttentionOld(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k)
        self.w_ks = nn.Linear(d_model, n_head * d_k)
        self.w_vs = nn.Linear(d_model, n_head * d_v)
        nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
        nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
        nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))

        self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5))
        self.layer_norm = nn.LayerNorm(d_model)

        self.fc = nn.Linear(n_head * d_v, d_model)
        nn.init.xavier_normal_(self.fc.weight)

        self.dropout = nn.Dropout(dropout)


    def forward(self, q, k, v, mask=None):

        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head

        sz_b, len_q, _ = q.size()
        sz_b, len_k, _ = k.size()
        sz_b, len_v, _ = v.size()

        residual = q

        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk
        k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk
        v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv

        mask = mask.repeat(n_head, 1, 1) # (n*b) x .. x ..
        output, attn = self.attention(q, k, v, mask=mask)

        output = output.view(n_head, sz_b, len_q, d_v)
        output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv)

        output = self.dropout(self.fc(output))
        output = self.layer_norm(output + residual)

        return output, attn
```

--------------------------------

### Independence of framework terminology: tile vs repeat

Source: https://github.com/arogozhnikov/einops/blob/main/README.md

Shows confusion between tile and repeat across frameworks, and how einops provides a uniform repeat notation.

```python
np.tile(image, (1, 2))    # in numpy
image.repeat(1, 2)        # pytorch's repeat ~ numpy's tile
```

```python
repeat(image, 'h w -> h (tile w)', tile=2)  # in numpy
repeat(image, 'h w -> h (tile w)', tile=2)  # in pytorch
repeat(image, 'h w -> h (tile w)', tile=2)  # in tf
repeat(image, 'h w -> h (tile w)', tile=2)  # in jax
repeat(image, 'h w -> h (tile w)', tile=2)  # in cupy
```

--------------------------------

### Import guess utility

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Imports a utility function 'guess' from a local module, used to check shape predictions in the tutorial.

```python
# utility to hide answers
from utils import guess
```

--------------------------------

### Space-to-depth with rearrange

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Rearranges spatial dimensions into channels, effectively performing space-to-depth with a 2x2 block.

```python
y = rearrange(x, "b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
guess(y.shape)
```

--------------------------------

### Display image grid using torchvision.utils.make_grid

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Uses torchvision.utils.make_grid to display the first 64 images from the fake batch with padding and normalization, then shows the result with matplotlib.

```python
#right
device = 'cpu'
plt.imshow(np.transpose(vutils.make_grid(fake_batch.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
```

```text
Result:
<matplotlib.image.AxesImage at 0x7fdd8d606908>
```

```text
Result:
<Figure size 432x288 with 1 Axes>
```

--------------------------------

### Combine composition and decomposition (b1 with height, b2 with width)

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Combines composition and decomposition: decomposes batch into b1 and b2, then composes b1 with height and b2 with width.

```python
# finally, combine composition and decomposition:
rearrange(ims, "(b1 b2) h w c -> (b1 h) (b2 w) c ", b1=2)
```

--------------------------------

### Alternative split for gated activation - channels with step

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Alternative split that takes channels with a step instead of consequent groups. Uses split=2.

```python
y1, y2 = rearrange(x, 'b (c split) h w -> split b c h w', split=2)
```

--------------------------------

### Rearrange 5D tensor and verify with parse_shape

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Shows how to rearrange a 5D tensor into a 4D form for 2D convolution, then use parse_shape to verify dimensions when restoring the original shape.

```python
# imagine we are working with 3d data
x_5d = rearrange(x, "b c x (y z) -> b c x y z", z=20)
# but we have only 2d convolutions.
# That's not a problem, since we can apply
y = rearrange(x_5d, "b c x y z -> (b z) c x y")
y = convolve_2d(y)
# not just specifies additional information, but verifies that all dimensions match
y = rearrange(y, "(b z) c x y -> b c x y z", **parse_shape(x_5d, "b c x y z"))
```

--------------------------------

### Benchmark output

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Measured timings for the old and new implementations. Old: 15.9 ms, new: 14.5 ms per loop.

```text
Output:
15.9 ms ± 990 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
14.5 ms ± 81.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

--------------------------------

### Import EinMix and torch modules

Source: https://github.com/arogozhnikov/einops/blob/main/docs/3-einmix-layer.ipynb

Import EinMix from einops.layers.torch and torch modules. The tutorial uses torch, but EinMix is available for other frameworks too.

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

# tutorial uses torch. EinMix is available for other frameworks too
from torch import nn
from torch.nn import functional as F
```

--------------------------------

### Compose batch and width axes

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Composes the batch and width axes into a single width dimension.

```python
# or compose a new dimension of batch and width
rearrange(ims, "b h w c -> h (b w) c")
```

--------------------------------

### Rearrange with b2=3

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Rearranges images with b2=3 to interweave along batch and spatial dimensions.

```python
rearrange(ims, "(b1 b2) (h1 h2) (w1 w2) c -> (h1 b1 h2) (w1 b2 w2) c", h1=3, w1=3, b2=3)
```

--------------------------------

### Benchmark channel shuffle functions

Source: https://github.com/arogozhnikov/einops/blob/main/scripts/pytorch_examples_source/Pytorch.ipynb

Times both channel shuffle implementations on GPU with %timeit and synchronizes CUDA.

```python
%timeit -n100 channel_shuffle_old(x, 8); torch.cuda.synchronize()
%timeit -n100 channel_shuffle_new(x, 8); torch.cuda.synchronize()
```

--------------------------------

### Depth-to-space with rearrange

Source: https://github.com/arogozhnikov/einops/blob/main/docs/2-einops-for-deep-learning.ipynb

Reverse of the previous space-to-depth operation, rearranging channels back into spatial dimensions.

```python
y = rearrange(x, "b (h1 w1 c) h w -> b c (h h1) (w w1)", h1=2, w1=2)
guess(y.shape)
```

--------------------------------

### Rearrange: Stacking List of Arrays

Source: https://github.com/arogozhnikov/einops/blob/main/docs/1-einops-basics.ipynb

Shows how rearrange can handle lists of arrays with the same shape. The list axis becomes the first axis (b) and is left in place.

```python
# rearrange can also take care of lists of arrays with the same shape
x = list(ims)
print(type(x), "with", len(x), "tensors of shape", x[0].shape)
# that's how we can stack inputs
# "list axis" becomes first ("b" in this case), and we left it there
rearrange(x, "b h w c -> b h w c").shape
```

```text
Output:
<class 'list'> with 6 tensors of shape (96, 96, 3)

```

```text
Result:
(6, 96, 96, 3)
```

--------------------------------

### Verify detection loss implementations match

Source: https://github.com/arogozhnikov/einops/blob/main/docs/4-pack-and-unpack.ipynb

Checks that the manual and einops-based detection loss functions produce identical results on random input.

```python
# check that results are identical
import torch
dims = dict(mask_h=6, mask_w=8, n_classes=19)
model_output = torch.randn([3, 5, 7, 5 + dims["mask_h"] * dims["mask_w"] + dims["n_classes"]])
for a, b in zip(loss_detection(model_output, **dims), loss_detection_einops(model_output, **dims)):
    assert torch.allclose(a, b)
```
Tokens:
Raw