Python - Einops

Table of Contents

Einstein-inspired notation for operations (Einops1) is a python package which provides extremely flexible and powerful opertions on tensors, e.g., permuting/reordering, reshaping (combining and partitioning along axes), reducing, repeating. Moreover, it supports numpy, tensorflow, pytorch, etc.

Installation

Einops can be directly installed from pypi. 

pip install einops

Usage

Basic functions

Basically, Einops provides 3 Functions, rearrange, reduce, repeat, which can be imported by 

from einops import rearrange, reduce, repeat

rearrange(tensor, pattern: str, **axes_lengths)
reduce(tensor: ~Tensor, pattern: str, reduction: Union[str, Callable[[~Tensor, List[int]], ~Tensor]], **axes_lengths: int) -> ~Tensor
repeat(tensor: ~Tensor, pattern: str, **axes_lengths) -> ~Tensor
  • tensor is a tensor (e.g., numpy.ndarray, tensorflow, pytorch) or a list of tensors with the same type and shape.
  • pattern is a string indicating rearrangement/reduction pattern.
  • reduction can be one of 'max', 'min', 'mean', 'sum', 'prod'.
  • axes_lengths are dimension related specifications.

rearrange

rearrange does not change the number of elements inside, and can achieve the operations like transpose, reshape, stack, squeeze, unsqueeze, concatenate, etc.

reduce

reduce combines elements in one or more axes. So it reduces the total number of elements, but keeps the ordering.

repeat

repeat produces repetition and tiling.

Layers

Einops also provides two corresponding layers, which can be directly integrated into any neural network. 

from einops.layers.chainer import Rearrange, Reduce
from einops.layers.gluon import Rearrange, Reduce
from einops.layers.keras import Rearrange, Reduce
from einops.layers.torch import Rearrange, Reduce
from einops.layers.tensorflow import Rearrange, Reduce

Taking pytorch for instace, a Rearrange layer or a Reduce layer can be generated by 

Rearrange(pattern, **axes_lengths)
Reduce(pattern, reduction, **axes_lengths)

Footnotes:

1

https://github.com/arogozhnikov/einops