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 einopsUsage
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) -> ~Tensortensoris a tensor (e.g.,numpy.ndarray,tensorflow,pytorch) or a list of tensors with the same type and shape.patternis a string indicating rearrangement/reduction pattern.reductioncan be one of'max','min','mean','sum','prod'.axes_lengthsare 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, ReduceTaking pytorch for instace, a Rearrange layer or a Reduce layer can be generated by
Rearrange(pattern, **axes_lengths)
Reduce(pattern, reduction, **axes_lengths)