einops

https://einops.rocks
Admin
einops is a flexible tensor operations library using Einstein-inspired notation that works with...
Tokens:40,498
Snippets:293
Trust Score:8.3
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install einops and numpy

Source: https://einops.rocks/4-pack-and-unpack

Install the necessary libraries to run the einops examples.

```bash
%pip install numpy einops -q
```

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

### Install Einops

Source: https://einops.rocks/

Install the einops library using pip. `uv` can also be used for installation.

```bash
pip install einops
```

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

### Additional einsum Usage Examples

Source: https://einops.rocks/api/einsum

Examples covering image filtering, matrix multiplication with unknown shapes, and matrix trace.

```python
# Filter a set of images:
>>> batched_images = np.random.randn(128, 16, 16)
>>> filters = np.random.randn(16, 16, 30)
>>> result = einsum(batched_images, filters,
...                 "batch h w, h w channel -> batch channel")
>>> result.shape
(128, 30)

# Matrix multiplication, with an unknown input shape:
>>> batch_shape = (50, 30)
>>> data = np.random.randn(*batch_shape, 20)
>>> weights = np.random.randn(10, 20)
>>> result = einsum(weights, data,
...                 "out_dim in_dim, ... in_dim -> ... out_dim")
>>> result.shape
(50, 30, 10)

# Matrix trace on a single tensor:
>>> matrix = np.random.randn(10, 10)
>>> result = einsum(matrix, "i i ->")
>>> result.shape
()
```

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

### Usage examples for einops.repeat

Source: https://einops.rocks/api/repeat

Demonstrates various tensor manipulation patterns including channel expansion, axis repetition, and upsampling.

```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)
```

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

### Running Einops Tests

Source: https://einops.rocks/

Command to install einops and pytest, and then run tests for specified frameworks like numpy, pytorch, and jax. The --pip-install flag is used to install dependencies in the current virtual environment.

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

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

### Einops Reduce Examples

Source: https://einops.rocks/api/reduce

This section provides various examples of using the einops.reduce function for different tensor operations.

```APIDOC
## Einops Reduce Function

### Description
Combines rearrangement and reduction using reader-friendly notation.

### Method
`reduce(tensor, pattern, reduction, **axes_lengths)`

### Endpoint
N/A (Python function)

### Parameters
- **tensor** (numpy.ndarray or similar) - The input tensor to be reduced.
- **pattern** (string) - A string defining the input and output tensor dimensions for rearrangement and reduction.
- **reduction** (string) - The type of reduction to perform (e.g., 'max', 'mean', 'sum').
- **axes_lengths** (dict, optional) - Keyword arguments specifying lengths for anonymous or named axes.

### Request Example
```python
import numpy as np
from einops import reduce

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

# Perform max-reduction on the first axis
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')

# Example with 4D tensor (batch of images)
x = np.random.randn(10, 20, 30, 40)

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

# Using anonymous axes for reduction
y1 = reduce(x, 'b c (h1 2) (w1 2) -> b c h1 w1', 'max')

# Adaptive 2D max-pooling to 3 * 4 grid
shape = reduce(x, 'b c (h1 h2) (w1 w2) -> b c h1 w1', 'max', h1=3, w1=4).shape

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

# Subtracting mean over batch for each channel
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')
```

### Response
#### Success Response (Output Tensor)
- The output is a tensor with dimensions rearranged and reduced according to the specified pattern and reduction method.

#### Response Example
```python
# Example output shape after global average pooling
(10, 20)
```

```

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

### Install Einops and Dependencies

Source: https://einops.rocks/1-einops-basics

Installs the einops library along with numpy and pillow for image manipulation. A kernel restart might be necessary after installation.

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

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

### Setup NumPy and Image Display

Source: https://einops.rocks/1-einops-basics

Imports NumPy and a utility function to display NumPy arrays as images in environments like Jupyter. This is useful for visualizing image data.

```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()
```

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

### Perform reduction operations with einops.reduce

Source: https://einops.rocks/api/reduce

Examples demonstrating various reduction patterns, including global pooling, max-pooling, and broadcasting subtractions.

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

```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')

# let's pretend now that x is a batch of images
# with 4 dims: batch=10, height=20, width=30, channel=40
>>> 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')

# 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 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')
```

```python
def reduce(tensor: Union[Tensor, list[Tensor]], pattern: str, reduction: Reduction, **axes_lengths: Size) -> Tensor:
    """
    einops.reduce combines rearrangement and reduction using reader-friendly notation.

    Some examples:

    ```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')

    # let's pretend now that x is a batch of images
    # with 4 dims: batch=10, height=20, width=30, channel=40
    >>> 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')

    # 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 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')

    ```

    Parameters:
        tensor: tensor of any supported library (e.g. numpy.ndarray, tensorflow, pytorch).
            list of tensors is also accepted, those should be of the same type and shape
        pattern: string, reduction pattern
        reduction: one of available reductions ('min', 'max', 'sum', 'mean', 'prod', 'any', 'all').
            Alternatively, a callable f(tensor, reduced_axes) -> tensor can be provided.
            This allows using various reductions like: np.max, np.nanmean, tf.reduce_logsumexp, torch.var, etc.
        axes_lengths: any additional specifications for dimensions

    Returns:
        tensor of the same type as input
    """
    try:
        if isinstance(tensor, list):
            if len(tensor) == 0:
                raise TypeError("Rearrange/Reduce/Repeat can't be applied to an empty list")
            backend = get_backend(tensor[0])
            tensor = backend.stack_on_zeroth_dimension(tensor)
        else:
            backend = get_backend(tensor)

        hashable_axes_lengths = tuple(axes_lengths.items())
        shape = backend.shape(tensor)
        recipe = _prepare_transformation_recipe(pattern, reduction, axes_names=tuple(axes_lengths), ndim=len(shape))
        return _apply_recipe(
            backend, recipe, cast(Tensor, tensor), reduction_type=reduction, axes_lengths=hashable_axes_lengths
        )
    except EinopsError as e:
        message = f' Error while processing {reduction}-reduction pattern "{pattern}".'
        if not isinstance(tensor, list):
            message += f"\n Input tensor shape: {shape}. "
        else:
            message += "\n Input is list. "
        message += f"Additional info: {axes_lengths}."
        raise EinopsError(message + f"\n {e}") from None
```

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

### Unpacking Reinforcement Learning Predictions

Source: https://einops.rocks/4-pack-and-unpack

Example of unpacking multiple outputs for reinforcement learning using einops.

```python
action_logits, reward_expectation, q_values, expected_entropy_after_action = \
    unpack(predictions_btc, [[n_actions], [], [n_actions], [n_actions]], 'b step *')
```

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

### Build PyTorch Model with Einops Reduce Layer

Source: https://einops.rocks/2-einops-for-deep-learning

This example shows how to build a PyTorch Sequential model using einops.layers.torch.Reduce for combined pooling and flattening. It's useful for simplifying complex tensor manipulations within a model architecture.

```python
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)
)
```

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

### Tensor Rearrangement Examples

Source: https://einops.rocks/api/rearrange

Demonstrates various tensor manipulation patterns including stacking, axis reordering, concatenation, flattening, and splitting using the rearrange function.

```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)
```

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

### Reduce Operations

Source: https://einops.rocks/1-einops-basics

Examples of reducing tensor dimensions using einops.reduce.

```python
x.mean(-1)
```

```python
reduce(x, 'b h w c -> b h w', 'mean')
```

```python
# average over batch
reduce(ims, "b h w c -> h w c", "mean")
```

```python
# the previous is identical to familiar:
ims.mean(axis=0)
# but is so much more readable
```

```python
# Example of reducing of several axes
# besides mean, there are also min, max, sum, prod
reduce(ims, "b h w c -> h w", "min")
```

```python
# this is mean-pooling with 2x2 kernel
# image is split into 2x2 patches, each patch is averaged
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "mean", h2=2, w2=2)
```

```python
# max-pooling is similar
# result is not as smooth as for mean-pooling
reduce(ims, "b (h h2) (w w2) c -> h (b w) c", "max", h2=2, w2=2)
```

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

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

### Einops Layers for PyTorch Models

Source: https://einops.rocks/

Integrate einops functionality directly into PyTorch models using `Rearrange` and `Reduce` layers. This example shows how `Rearrange` can replace manual flattening.

```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),
)
```

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

### Width-to-Height Tensor Reshaping with Einops

Source: https://einops.rocks/

Example of reshaping a tensor to change its width and height dimensions using einops rearrange, demonstrating flexibility beyond standard operations.

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

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

### Select framework backend

Source: https://einops.rocks/2-einops-for-deep-learning

Define the target framework for the tutorial.

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

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

### Initialize framework backend

Source: https://einops.rocks/2-einops-for-deep-learning

Configure the environment for either TensorFlow or PyTorch.

```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
```

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

### Serving Documentation Locally

Source: https://einops.rocks/

Command to build and serve the project's documentation locally using hatch, typically accessible at http://localhost:8000/.

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

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

### Initialize EinMix

Source: https://einops.rocks/3-einmix-layer

Import the EinMix layer and necessary PyTorch modules.

```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
```

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

### Decompose and Compose Axes

Source: https://einops.rocks/1-einops-basics

Examples of splitting and merging tensor axes using rearrange.

```python
rearrange(ims, "(b1 b2) h w c -> b1 b2 h w c ", b1=2).shape
```

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

```python
# slightly different composition: b1 is merged with width, b2 with height
# ... so letters are ordered by w then by h
rearrange(ims, "(b1 b2) h w c -> (b2 h) (b1 w) c ", b1=2)
```

```python
# move part of width dimension to height.
# we should call this width-to-height as image width shrunk by 2 and height doubled.
# but all pixels are the same!
# Can you write reverse operation (height-to-width)?
rearrange(ims, "b h (w w2) c -> (h w2) (b w) c", w2=2)
```

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

### Framework adaptation

Source: https://einops.rocks/1-einops-basics

Demonstrates that einops operations are universal across different deep learning frameworks.

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

```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')
```

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

### Rearrange Tensor Shape

Source: https://einops.rocks/1-einops-basics

Rearrange the dimensions of a tensor. This example shows how a list of tensors is treated as a batch and rearranged.

```python
rearrange(x, "b h w c -> b h w c").shape
```

```python
rearrange(x, "b h w c -> h w c b").shape
```

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

### Einops Import for Array API Standard

Source: https://einops.rocks/

Demonstrates how to import einops rearrange, showing the recommended import path for compatibility with the Python Array API standard.

```python
from einops import rearrange
=> from einops.array_api import rearrange
```

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

### Manual unpacking strategies

Source: https://einops.rocks/4-pack-and-unpack

Demonstrate different ways to split tensors using custom shape specifications.

```python
# simple unpack by splitting the axis. Results are (h, w, 3) and (h, w, 1)
rgb, depth = unpack(image_rgbd, [[3], [1]], "h w *")
# different split, both outputs have shape (h, w, 2)
rg, bd = unpack(image_rgbd, [[2], [2]], "h w *")
# unpack to 4 tensors of shape (h, w). More like 'unstack over last axis'
[r, g, b, d] = unpack(image_rgbd, [[], [], [], []], "h w *")
```

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

### Interweave Pixels of Different Images

Source: https://einops.rocks/1-einops-basics

Interweave pixels from different images by rearranging dimensions. This example shows how to combine images based on batch dimensions.

```python
rearrange(ims, "(b1 b2) h w c -> (h b1) (w b2) c ", b1=2)
```

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

### Perform reductions with einops.reduce

Source: https://einops.rocks/api/reduce

Demonstrates various reduction patterns including max-pooling, global average pooling, and per-channel mean subtraction.

```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')

# let's pretend now that x is a batch of images
# with 4 dims: batch=10, height=20, width=30, channel=40
>>> 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')

# 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 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')
```

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

### Reduce Tensor with Max Operation

Source: https://einops.rocks/1-einops-basics

Perform a reduction operation on a tensor. This example computes the maximum value in each image and then shows the difference from the original image.

```python
x = reduce(ims, "b h w c -> b () () c", "max") - ims
rearrange(x, "b h w c -> h (b w) c")
```

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

### Implement Channel Shuffle

Source: https://einops.rocks/pytorch-examples.html

Demonstrates the reduction of complex view and transpose operations into a single einops rearrange call for channel shuffling.

```python
def channel_shuffle_old(x, groups):
    batchsize, num_channels, height, width = x.data.size()

    channels_per_group = num_channels // groups
    
    # reshape
    x = x.view(batchsize, groups, 
        channels_per_group, height, width)

    # transpose
    # - contiguous() required if transpose() is used before view().
    #   See https://github.com/pytorch/pytorch/issues/764
    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batchsize, -1, height, width)

    return x
```

```python
def channel_shuffle_new(x, groups):
    return rearrange(x, 'b (c1 c2) h w -> b (c2 c1) h w', c1=groups)
```

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

### Perform backpropagation

Source: https://einops.rocks/2-einops-for-deep-learning

Demonstrate gradient calculation through Einops operations.

```python
y0 = x
y1 = reduce(y0, "b c h w -> b c", "max")
y2 = rearrange(y1, "b c -> c b")
y3 = reduce(y2, "c b -> ", "sum")

if flavour == "tensorflow":
    print(reduce(tape.gradient(y3, x), "b c h w -> ", "sum"))
else:
    y3.backward()
    print(reduce(x.grad, "b c h w -> ", "sum"))
```

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

### Pack and unpack GAN inputs

Source: https://einops.rocks/4-pack-and-unpack

Shows how to pack true and fake images for model processing and unpack the results.

```python
input_ims, ps = pack([true_images, fake_images], '* h w c')
true_pred, fake_pred = unpack(model(input_ims), ps, '* c')
```

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

### Compare Linear Layer Implementations

Source: https://einops.rocks/3-einmix-layer

Comparison between standard torch.einsum and the EinMix layer implementation.

```python
weight = <...create and initialize parameter...>
bias = <...create and initialize parameter...>
result = torch.einsum('tbc,cd->tbd', embeddings, weight) + bias
```

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

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

### Use framework-specific layers

Source: https://einops.rocks/2-einops-for-deep-learning

Import and initialize einops layers for various deep learning frameworks.

```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
```

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

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

### Perform Tensor Products with Named Axes

Source: https://einops.rocks/api/einsum

Demonstrates the basic usage of einsum with multiple tensors and a pattern string.

```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")
```

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

### Visualize generative model image batches

Source: https://einops.rocks/pytorch-examples.html

Methods for creating image grids from model output batches using torchvision utilities or einops.

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

```python
padded = F.pad(fake_batch[:64], [1, 1, 1, 1])
plt.imshow(rearrange(padded, '(b1 b2) c h w -> (b1 h) (b2 w) c', b1=8).cpu())
```

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

### Importing Dependencies

Source: https://einops.rocks/pytorch-examples.html

Standard imports required for PyTorch and einops operations.

```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
```

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

### Pack multi-modal tokens

Source: https://einops.rocks/4-pack-and-unpack

Demonstrates packing tokens from various sources into a single tensor.

```python
all_inputs = [text_tokens_btc, image_bhwc, task_token_bc, static_tokens_bnc]
inputs_packed, ps = pack(all_inputs, 'b * c')
```

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

### Verification of Detection Loss Implementation

Source: https://einops.rocks/4-pack-and-unpack

Validates that the manual and einops implementations produce identical results using torch.

```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)
```

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

### Import Einops functions

Source: https://einops.rocks/2-einops-for-deep-learning

Initial imports for rearranging and reducing tensors.

```python
from einops import rearrange, reduce
```

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

### Prepare Data for Vision Transformer

Source: https://einops.rocks/4-pack-and-unpack

Initializes random patch tokens and zero-initialized class tokens, representing preparations for a Vision Transformer model.

```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])
```

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

### Einops Einsum for Generalized Dot Products

Source: https://einops.rocks/

Utilize `einsum` for generalized dot products with support for multi-lettered axes and a pattern argument that goes last. It is compatible with multiple frameworks.

```python
from einops import einsum, pack, unpack
# einsum is like ... einsum, generic and flexible dot-product
# but 1) axes can be multi-lettered  2) pattern goes last 3) works with multiple frameworks
C = einsum(A, B, 'b t1 head c, b t2 head c -> b head t1 t2')
```

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

### Implement ViT with standard NumPy

Source: https://einops.rocks/4-pack-and-unpack

A manual implementation of the ViT token processing pipeline using standard NumPy operations for comparison.

```python
def vit_vanilla(class_tokens, patch_tokens):
    b, h, w, c = patch_tokens.shape
    class_tokens_b1c = class_tokens[:, np.newaxis, :]
    patch_tokens_btc = np.reshape(patch_tokens, [b, -1, c])
    input_packed = np.concatenate([class_tokens_b1c, patch_tokens_btc], axis=1)
    output_packed = transformer_mock(input_packed)
    class_token_emb = np.squeeze(output_packed[:, :1, :], 1)
    patch_tokens_emb = np.reshape(output_packed[:, 1:, :], [b, h, w, -1])
    return class_token_emb, patch_tokens_emb

class_token_emb2, patch_tokens_emb2 = vit_vanilla(class_tokens, patch_tokens)
assert np.allclose(class_token_emb, class_token_emb2)
assert np.allclose(patch_tokens_emb, patch_tokens_emb2)
```

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

### ShuffleUnitNew Module for ShuffleNet

Source: https://einops.rocks/pytorch-examples.html

A revised implementation of the ShuffleUnit, potentially using more modern PyTorch features like `Rearrange`. This snippet is incomplete, showing only the `__init__` method.

```python
class ShuffleUnitNew(nn.Module):
    def __init__(self, in_channels, out_channels, groups=3, 
                 grouped_conv=True, combine='add'):
        super().__init__()
        first_1x1_groups = groups if grouped_conv else 1
        bottleneck_channels = out_channels // 4
        self.combine = combine
        if combine == 'add':
            # ShuffleUnit Figure 2b
            self.left = Rearrange('...->...') # identity
            depthwise_stride = 1
        else:
            # ShuffleUnit Figure 2c
            self.left = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
            depthwise_stride = 2
            # ensure output of concat has the same channels as original output channels.
            out_channels -= in_channels
            assert out_channels > 0

        self.right = nn.Sequential(

```

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

### Depth-to-Space Operation with Einops Rearrange

Source: https://einops.rocks/

Demonstrates two ways to implement the depth-to-space operation using einops rearrange, highlighting the ability to precisely define the output tensor structure.

```python
# depth-to-space
rearrange(x, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=2, w2=2)
```

```python
rearrange(x, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=2, w2=2)
```

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

### Import utility

Source: https://einops.rocks/2-einops-for-deep-learning

Import a helper function for hiding answers.

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

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

### Refactoring Super-Resolution Network

Source: https://einops.rocks/pytorch-examples.html

Replacing custom PixelShuffle logic with einops Rearrange for better framework portability and memory efficiency.

```python
class SuperResolutionNetOld(nn.Module):
    def __init__(self, upscale_factor):
        super(SuperResolutionNetOld, self).__init__()

        self.relu = nn.ReLU()
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x
```

```python
def SuperResolutionNetNew(upscale_factor):
    return nn.Sequential(
        nn.Conv2d(1, 64, kernel_size=5, padding=2),
        nn.ReLU(inplace=True),
        nn.Conv2d(64, 64, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(64, 32, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(32, upscale_factor ** 2, kernel_size=3, padding=1),
        Rearrange('b (h2 w2) h w -> b (h h2) (w w2)', h2=upscale_factor, w2=upscale_factor),
    )
```

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

### Implement ViT with Einops

Source: https://einops.rocks/4-pack-and-unpack

Uses pack and unpack to handle class and patch tokens in a transformer pipeline.

```python
def vit_einops(class_tokens, patch_tokens):
    input_packed, ps = pack([class_tokens, patch_tokens], "b * c")
    output_packed = transformer_mock(input_packed)
    return unpack(output_packed, ps, "b * c_out")

class_token_emb, patch_tokens_emb = vit_einops(class_tokens, patch_tokens)

class_token_emb.shape, patch_tokens_emb.shape
```

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

### Display First Image

Source: https://einops.rocks/1-einops-basics

Displays the first image from the loaded dataset. Note that the entire 4D tensor cannot be rendered as a single image.

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

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

### Perform Pooling Operations

Source: https://einops.rocks/2-einops-for-deep-learning

Demonstrates 1D and volumetric pooling using the reduce function with einops expressions.

```python
reduce(x, '(t 2) b c -> t b c', reduction='max')
```

```python
reduce(x, 'b c (x 2) (y 2) (z 2) -> b c x y z', reduction='max')
```

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

### Refactoring ConvNet Implementation

Source: https://einops.rocks/pytorch-examples.html

Comparison between a manual class-based ConvNet and a sequential implementation using einops for shape management.

```python
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

conv_net_old = Net()
```

```python
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)
)
```

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

### Implement TokenMixer with EinMix

Source: https://einops.rocks/3-einmix-layer

Creates a sequential TokenMixer module using EinMix to handle axis mixing. Note that this implementation intentionally excludes residual connections.

```python
def TokenMixer(num_features: int, n_patches: int, expansion_factor: int, dropout: float):
    n_hidden = n_patches * expansion_factor
    return nn.Sequential(
        nn.LayerNorm(num_features),
        Mix("b hw c -> b hid c", weight_shape="hw hid", bias_shape="hid", hw=n_patches, hidden=n_hidden),
        nn.GELU(),
        nn.Dropout(dropout),
        Mix("b hid c -> b hw c", weight_shape="hid hw", bias_shape="hw",  hw=n_patches, hidden=n_hidden),
        nn.Dropout(dropout),
    )
```

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

### Reduce and Repeat as Inverse Operations

Source: https://einops.rocks/1-einops-basics

Demonstrates reduce and repeat as inverse operations. First, an image is repeated, then reduced using 'min' to return the original tensor.

```python
repeated = repeat(ims, "b h w c -> b h new_axis w c", new_axis=2)
reduced = reduce(repeated, "b h new_axis w c -> b h w c", "min")
assert numpy.array_equal(ims, reduced)
```

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

### Swap Height and Width with Verbose Names

Source: https://einops.rocks/1-einops-basics

Demonstrates using more descriptive names for axes ('height', 'width', 'color') in the `rearrange` operation, achieving the same result as using short names.

```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
```

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

### Import numpy

Source: https://einops.rocks/4-pack-and-unpack

Import numpy for data manipulation.

```python
import numpy as np
```

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

### Refactoring Gram Matrix Calculation

Source: https://einops.rocks/pytorch-examples.html

Using torch.einsum to simplify Gram matrix computation compared to manual transpose and batch matrix multiplication.

```python
def gram_matrix_old(y):
    (b, ch, h, w) = y.size()
    features = y.view(b, ch, w * h)
    features_t = features.transpose(1, 2)
    gram = features.bmm(features_t) / (ch * h * w)
    return gram
```

```python
def gram_matrix_new(y):
    b, ch, h, w = y.shape
    return torch.einsum('bchw,bdhw->bcd', [y, y]) / (h * w)
```

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

### Display Second Image

Source: https://einops.rocks/1-einops-basics

Displays the second image from the loaded dataset.

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

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

### Refactor FastText Pooling

Source: https://einops.rocks/pytorch-examples.html

Demonstrates replacing manual tensor permutation and pooling with a sequential model using Einops operations.

```python
class FastTextOld(nn.Module):
    def __init__(self, vocab_size, embedding_dim, output_dim):
        super().__init__()
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.fc = nn.Linear(embedding_dim, output_dim)
        
    def forward(self, x):
        
        #x = [sent len, batch size]
        
        embedded = self.embedding(x)
                
        #embedded = [sent len, batch size, emb dim]
        
        embedded = embedded.permute(1, 0, 2)
        
        #embedded = [batch size, sent len, emb dim]
        
        pooled = F.avg_pool2d(embedded, (embedded.shape[1], 1)).squeeze(1) 
        
        #pooled = [batch size, embedding_dim]
                
        return self.fc(pooled)
```

```python
def FastTextNew(vocab_size, embedding_dim, output_dim):
    return nn.Sequential(
        Rearrange('t b -> t b'),
        nn.Embedding(vocab_size, embedding_dim),
        Reduce('t b c -> b c', 'mean'),
        nn.Linear(embedding_dim, output_dim),
        Rearrange('b c -> b c'),
    )
```

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

### Implement MLPMixer TokenMixer

Source: https://einops.rocks/3-einmix-layer

Standard PyTorch implementation of the MLP and TokenMixer components from MLPMixer.

```python
class MLP(nn.Module):
    def __init__(self, num_features, expansion_factor, dropout):
        super().__init__()
        num_hidden = num_features * expansion_factor
        self.fc1 = nn.Linear(num_features, num_hidden)
        self.dropout1 = nn.Dropout(dropout)
        self.fc2 = nn.Linear(num_hidden, num_features)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x):
        x = self.dropout1(F.gelu(self.fc1(x)))
        x = self.dropout2(self.fc2(x))
        return x


class TokenMixer(nn.Module):
    def __init__(self, num_features, num_patches, expansion_factor, dropout):
        super().__init__()
        self.norm = nn.LayerNorm(num_features)
        self.mlp = MLP(num_patches, expansion_factor, dropout)

    def forward(self, x):
        # x.shape == (batch_size, num_patches, num_features)
        residual = x
        x = self.norm(x)
        x = x.transpose(1, 2)
        # x.shape == (batch_size, num_features, num_patches)
        x = self.mlp(x)
        x = x.transpose(1, 2)
        # x.shape == (batch_size, num_patches, num_features)
        out = x + residual
        return out
```
Tokens:
Raw