Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation
Authors: Alex Rogozhnikov
Venue: ICLR 2022 Oral (2022)
Keywords: tensor manipulations
tensor transformation
einops
einstein notation
einsum

Abstract: Tensor computations underlie modern scientific computing and deep learning.
A number of tensor frameworks emerged varying in execution model, hardware support, memory management, model definition, etc.
However, tensor operations in all frameworks follow the same paradigm.
Recent neural network architectures demonstrate demand for higher expressiveness of tensor operations.
The current paradigm is not suited to write readable, reliable, or easy-to-modify code for multidimensional tensor manipulations. 
Moreover, some commonly used operations do not provide sufficient checks and can break a tensor structure.
These mistakes are elusive as no tools or tests can detect them.
Independently, API discrepancies complicate code transfer between frameworks.
We propose einops notation: a uniform and generic way to manipulate tensor structure, that significantly improves code readability and flexibility by focusing on the structure of input and output tensors.
We implement einops notation in a Python package that efficiently supports multiple widely used frameworks and provides framework-independent minimalist API for tensor manipulations.

Published as a conference paper at ICLR 2022
EINOPS: CLEAR AND RELIABLE TENSOR MANIPULATIONS WITH EINSTEIN-LIKE NOTATION
Alex Rogozhnikoy a alex. wowozhnikkooaa........mmmmm
ABSTRACT
Tensor computations underlie modern scientific computing and deep learning. A number of tensor frameworks emerged varying in execution model, hardware support, memory management, model definition, etc. However, tensor operations in all frameworks follow the same paradigm. Recent neural network architectures demonstrate demand for higher expressiveness of tensor operations. The current paradigm is not suited to write readable, reliable, or casy-to-modify code for mul- tidimensional tensor manipulations. Moreover, some commonly used operations do not provide sufficient checks and can break a tensor structure. These mistakes are elusive as tools tests detect them. Independently, API discrepancies complicate code transfer between frameworks. We propose einops notation: a uniform and generic way manipulate tensor structure, that significantly improves code readability and flexibility by focusing on the structure of input and output ten- sors. We implement einops notation Python package that efficiently supports multiple widely used frameworks and provides Tramewor--indepenee minimalist API for tensor manipulations.
INTRODUCTION
Deep learning (DL) over the past decade achieved a significant progress in analysis and synthesis of images, audio and text (Bengio et al., 2017; Aggarwal et al., 2018; Foster, 2019). Tools relying on these methods became a commodity in production data pipelines.
frameworks for DI., such as pytorch (Paszke et al., 2019), tensor- flow (Abadi et al., 2015), mxnet (Chen al., 2015), jax (Bradbury al., 2018), and others, vary in numerous aspects, but their core functionality is built around efficient computations on n-dimensional arrays (tensors for brevity ). API exposed by frameworks for tensor operations follow the same approach, that combines high efficiency (specialized hardware can be utilized) and user convenience: computations can be expressed in high-level languages, such as Python, using a limited number of exposed and usually pre-compiled operations.
Available assearch-and-producc
Due to growing usage of DI. in production and rising complexity models, becomes increasingly more important to provide programming interfaces that enable reliable and scalable development. We demonstrate that approach to define tensor operations taken by existing frameworks does not encourage writing code that is easy to interpret, maintain modify; additionally, some of the core operations do not conduct sufficient checks and can lead to hard-to-catch mistakes. To address these problems, we propose Einstein-like notation for operations, called einops. We implement this approach in a Python (Van Rossum & Drake, 2009) package to allow simple integration of notation into existing code across variety of frameworks.
Outline We first briefly describe mainstream approach for tensor operations and point to its issues with examples. We review previously proposed ideas resolve mentioned problems in related works. Our approach, einops a verbose notation for tensor manipulation - is introduced next, followed by the code examples case studies. We implement the notation in Python package, and explain main design choices in implementation details section. We conclude with discussion of einops role and common criticisms.
CCureetlyy at Herophilus, Inc. 'Ou sincere apologies readers with backgrounds in mathematies and physics for possible confusion.

---

Published as a conference paper at ICLR 2022
2 MAINSTREAM APPROACH FOR TENSOR OPERATIONS
Nowadays tensor programming is dominating in DL and playing a crucial role scientific computing. It first appeared in APL (Iverson, 1962), was popularized by MATLAB (Matlab, 1993) and was spread in Python community by numpy (Harris et al., 2020), and now is provided by multiple frameworks. Currently, tensor manipulations in all mainstream frameworks are based the following assumptions:
Tensor an n-dimensional array with shape (e.g. 6 3) and data type (e.g. float32).
When elements matched across tensors, are aligned by order. Conventions regulate cases of different dimensionalities, e.g. broadcasting in Numpy (2021).
Multiple operations act differently on axes. Those operations either have conventions about axes order (e.g. recurrent units, convolutions; convolutions expect input be either BHWC or BCHW ordered) or should be provided with indices axes that should be treated separately:"
x. max
# some operations are provided with indices
x.swap_ swaapp_axs((e axes(e, 1) # of specially treated axes
The mixed option is possible when operation has defaults for "special axes" that can be overridden during a call.
Benefits of the mainstream approach are simple API and a baseline implementation, as well as versatility: lot of research code operates solely using this kind of operations for number crunching. The drawbacks are absence of semantics in an operation and no way to reflect expected input and output. To predict output after chain of manipulations researcher/engineer has to carefully keep in the memory (or in the code comments) shape and contents of each intermediate tensor and thus keep track of every operation.
In this listing a batch of images is collapsed into a single image by placing images in a row next to each other
Lm_stack has shape [b, 3, h, w] y1 im_stack. transpose(2, 0, 3, 1). reshpppee(,, b w, 3) 1). reshapeeh,,,, b 3)
y2 im_stack. transpose(2,
One of y1, y2 contains correct image, while the other is irrelevant it requires some time figure out which is what. A typical computation chain contains multiple tensor operations (only two in this example) and requires writing down all intermediate steps "debug" the code. Annoyingly, resulting tensors y1 and y2 have the same shapes and mistake in the code may stay under the radar for a long time since the code never errors out. The lack of stronger checks is a weak point of conventional operations.
In most cases we also cannot meaningfully visualize intermediate layers, so there no way narrow down searches for a problem source. Thus, researcher/enginner has to check all the code after failure.
Traditional operations restrict code flexibility:: any change in the shape agreements between parts of the code is hard align and related code fragments (frequently located different places) should be updated simultaneously. In the next code fragment we add omitted batch dimension to the code from the vision permutator (Hou al., 2021), and then update code to support depth:
pytorch-like code without batch dimension, as in the paper x_h x.reshape(H, N, permute((, 1, reshepep, W, H+S) proj_h(x__]__ proj__(x__h.rreheeeee reshapee(,, W, H, s).permute(2, 1, 0, reshhappe,,, W, with batch dimension
x_h = x.reshape(B, H, W, N, >erpemuuee(o, 3, 2, 1, re.heppap) N, W, HAS)
x_h = proj_h(x_h) res.hapeees N, w, H, permuttte 3, 1, reshaape) H, W, C)
d with batch and depth dimension
x_h x.reshape(B, W, N, s).permute(o, 4, 2, 3, 1, reshapee) N, W, H*S) N, w, D, H, s). 4, 2, 3, 1, reshapee, H, W,
x_h
D, C)
PPseudocode in the paper corresponds to numpy unless otherwise stated. There is no way to write cross- framework code. This problem we partially address with proposed einops notation.
2

---

Published as a conference paper at ICLR 2022
Modifications very error-prone all indices should be recomputed and order of reshape operands should be verified. Uncommonly, this fragment quite delf-documnnting since final reshape in each line hints the intended order of axes after the transposition. In DL code it is common to use x. trannsspeee 3, 2, 2) expecting other users to recognize a familiar pattern or leaving a comment." Related, transposition in the code requires operating with three contexts in mind (original order of axes, permutation, result order), and even when simplified to just permutation, it's unclear if permutation (2 o 3 1) is inverse to 3 0 2).
Less critical, but still notable source of mistakes is O-based enumeration of axes (Julia, MATLAB and R use I-based), while we propose the framework not relying on axis numeration all. Common to refer to axis with index 0 as "first axis/dimension" (see e.g. numpy documentation), and there is no way to change this habit and avoid confusion hetween human communication and code.
Finally, common tensor operations can easily break the tensor structure. For example, reshape, common operation in the DL code, easily breaks the tensor structure because whole tensor is treated as sequence and no connection is assumed between axes in input and output, see Figure 1.
e
e e
(a) x shape = (92, 3, 96)
x. r45 3, -1) shape = (46, 3, 192)
(c) x. reshape(48, 3, -1) shape = (48, 3, 184)
Figure 1: Example of breaking the tensor structure in 3d tensor. We non-convential order HCW for visalizations: (a) original image; (b) partial mixing results in new colors; (c) mixing of all axes. Result shapes are shown underneath.
Aforementioned implementation mistakes result in rejecting valuable research or drawing incor- rect conclusion about the data. Moreover, there is no recipe to identify incorrect implementation, because the result tensor shape and data type are not affected by error can be obscured: poor performance can be misattributed to a number of other factors: data, hyperparameters, method, etc. All pointed issues have catastrophic importance for research costs and timeline in a setting where one experiment requires dozens to thousands of GPU-hours.
Even the presence of an
number of recently proposed architectures (Wu al., 2021; Tolstikhin al., 2021; Jumper et al., 2021; Liu et al., 2021; Hou et al., 2021; Touvron et al., 2021) demonstrate that conservative approach with fixed ordering of axes (like BCHW for convolutional networks) is not sufficient. Existing frameworks carry this imprint that does not help with searches for new architectures.
RELATED WORKS
A commonly taken approach to increase reliability is assigning names/labels to tensor dimensions (below refer to them labeled tensors). Most known implementations are xarray (Hoyer & Hamman, 2017), labeled tensor in tensorflow (Christiansen & Hoyer, 2016), namedtensor babeled_tensor
(2019), and named tensors in Pytorch (2019). Despite being around for many years, this idea got little adoption, and not used in the DI code: development stopped for labeled_ tensor and namedtensor, and named tensors in pytorch are still experimental.
In this approach, operations match tensor axes based on labels (common choice of label is string name), and rely on axis label instead of axes order, e.g.:
has (x, y, height) and has axes (time, x, x1 x2 # result has axes (x, y, height, time), maybe in a different order x1. mean( "highht) reduce over axis named 'height"
Out of 50 last usages of torch. permute in Python GitHub on 22 Nov 2021, only in 4 cases code had comments about result shape, in other information was provided in any other form, including reshape.
3

---

Published as a conference paper at ICLR 2022
Axes matching rules vary across implementations. However, describe common issues why this approach did not gain wide support in the DL community:
Operations focus modified axes, and neither describe input output; user has remember axes labels for each intermediate tensors.
Less control over data layout: order of axes may significantly influence speed (Weber & Goesele, 2017), but is not transparent.
Names should be strictly defined and mistakes in names or their alignment between modules may result in the wrong computations, not an exception, c.g. for namedtensor package:
has axes (height, width)
# x2 has axes (h, w)
# x3 has axes (heigt)
x2 x3 # has axes (height width
heigt) in some order
Adoption requires a strong buy-in decision, as all code should be axis label-centric. In contrary. transition labeled tensors in isolated code pieces only (e.g. functions) does not prevent (more elusive) errors in interaction between functions, but introduces constant conversions between the code styles. Third-party components need wrapping to support labels.
Symmetric and antisymmetric tensors with multiple identical axes pose an additional challenge: all axes labels should be unique to allow matching and axis specification.
Labeled tensors have issues with integration of DL. blocks (details in Appendix A).
Proposed implementations of labeled tensors also break the common principle of decomposition in software engineering, which states that every function has its own scope with input names and names for intermediate variables. Everything that is shared between scopes should be described in the function signature. Whenever an internal structure of passed or returned entity should be shared between scopes, appeccaas/inteefac//pooocs is introduced describe passed argument. However, the concept of labeled tensor breaks this logic: it is assumed that called and calling functions agree on axes names, but no way to document these expectations is proposed.
Alternative approach to increase readability and reliability of tensor-operating code is to deliberately set interface restrictions only on large neural modules such as language models, encoders or decoders, as in NeMo Kuchaiev et al. (2019). While allowing to reuse and wrap existing code to glue large components, this approach does not improve internals of the modules where problems are harder to locate and code is less documented. These improvements of high-level interfaces still have their challenges, for example language model can expect to manipulate sequences of letters and thus expects axis "letter". However, surrounding code may try phonemes. Thus, relabeling of axes may be required to "connect" subsystems.
for prediction of words, pixels
In 2011, M. Wiebe introduced operation einsum in numpy package. With some simplifications (absence of covariant and contravariant indices, contracted dimension may be not repeated) e einsum mimics Einstein summation rule commonly used in physics. numpy. numpy.einsum einsum stands out from the rest of numpy API and for a long time rarely was mentioned in tutorials. However, function universality and expressiveness were beneficial, and it was ported other packages: tensorflow, pytorch, mxnet, chainer Tokui et al. (2019), etc.
A, B) it matrix mul multiplication compy.einsumt 'ijk->ij', C) # sum over last axis numpy. numpy.einsum( 'ij,ji->', A,
einsum((
trace of matrix product
einops revisits and expands this approach with emphasis on complex tensor layouts and rear- rangements, additional checks and broader functionality. In our work we try to align interface with einsum allow smooth simultaneous usage. However, interface adjustments (such support for multi-character axes names) are necessary. Most of our changes can be readily applied to einsum. Detailed discussion of differences between einops and einsum is given in Appendix B.
There is an ongoing research to create languages for low-level definition of tensor operations with explicit indexing, e.g. Tensor Comprehensions (Vasilache et al., 2018) and Dex (Paszke et al., 2021).
4

---

Published as a conference paper at ICLR 2022
4 einops
Einstein- Inseiredd Notation for OPerationS, einops, is our proposal to address problems mentioned in Section 2. The core of our approach is notation for transformation patterns and, in its basic view, this notation enumerates the tensor elements in one-to-one-coree coro-ooooo-oeessoodddeeee to the set of axes. As seen in examples below, we allow number of axes to be different from tensor dimensionality. The notation uses simple syntax, formally defined in Appendix C, and is based on the following rules:
axis present only in the input (the left hand side) is reduced (e.g. with max-reduction)
axis present only in the output is "repeated" (tensor values are the same for all values of axes)
all axis identifiers either side of expression should be unique
Examples of einops transformation patterns are
"b >> b h C' # transpose
"b >> b c'
# reduce on h, w
"b b c w'
# repeat on h, w
"(h1 h2) (w1 w2) c (h1 w1) h2 w2 split image to patches, stack them
Each tensor pattern ensures one-to-one mapping between element's positions in the tensor and values of axis variables. This requires all axes used in tensor pattern be different (thus traces, permitted by einsum, are not provided by einops). This also requires that ellipsis can be used only once within tensor pattern.
The main novelty of our notation is the composition and decomposition of axes denoted by parenthesis. (De)composition uses C-ordering convention, intuitively associated with digits number:
x of shape (10, 10, 10, 10), then x[6, 2, 4, 9] y[6249] = rearrange(x, 'a
>> (a b d)') d)
Changing the rightmost of "digits" changes composed index in "small steps", while any change in leftmost results in "large steps", even when axes are not decimal:
# Rearrange pattern that composes 3 axes into one: i1 i2 is -> (i1 i2 i3)
# Let original array have shape of [2, 3, 2], result has a length 2x3X2=12
i1
1
i2
8 1
2
8
1
2
i3
1 1
8
9
0
1 1
final position
0
1
3 4
5
6
8
10
11
Reverse pattern (in i2 i3) i1 i2 i3 uses the same bijection to decompose an axis into three. Since axis can decomposed in multiple ways (e.g. 12 could be represented as 2 x x or 1 x 12 x 1 or 3 1 x 4, etc.), additional axis size specifications are required during decomposition. The following rule is helpful for C-ordered arrays (default ordering in all current backends): in case the order of axes does not change, result of rearrange is still view of the original data. For example, rearrangement (a b c) (d f) (g h) > a (c d) e (f gigh h)' returns a view and no copy is required.
Axes can be referred to by their size. These anonymous axes have the same role as named, but due to the lack of name they can't be matched across different tensors. Unitary axes have special meaning- they do not correspond to an axis variable, and thus their behavior is separate from anonymous axes.
'h >> h 1' " h w
# add unitary axis
repeat values along new anonymous axis
1 >> w²
# remove unitary axis and reduce on anonymous axis length
h w +> W h
transpose two last dimensions
<> (...)
b c² compose all but first and last dimensions
and move resulting new axis front
b
# reduce on all but first two dimensions
All anonymous axes are treated as different even if they share length. Similar to einsum, ellipsis works wildcard for
axes, their number and lengths are derived from input shape(s).
5

---

Published as a conference paper at ICLR 2022
Ellipses are not named (grammar allows adding names later), thus all ellipses within transformation pattern refer to the same group of unnamed axes.
Pattern composition and decomposition became particularly helpful to leverage existing operations for data of higher dimensionality, E.g. if an attention function accepts tensors k, 9, v of shape [batch, turn into multi-head attention for 3-dimensional data by composing height,
seq, channel], 1, width and depth to a single dimension, and grouping head and batch dimension to ensure independent processing of attention heads: blocks can be "upgraded" by rearranging inputs and outputs.
w d head -> head) w d) c. Likewise, other neural
einops provides three functions (rearrange, reduce, and repeat) which are shown below in examples (additional axes specifications are provided with **kwargs):
organize 16 images in 4x4 grid
rearrange perrangeeimm, im, '(b1 b2) h w c (b1 h) (b2 w) c', b1=4, b2=4)
max max-pooling with kernel of size 2x2
reduce im, 'b c (h h2) (w w2). -> b c. w', 'max', h2=2, w2=2)
2x upsampling of individual image by repeating pixels repeet lim, 'h
(h h2) (w w2) c', h2=2, w2=2, c=3)
While all patterns could be handled by a single function instead of three, we made an explicit choice to separate scenarios of "adding dimensions" (repeat), "removing dimensions" (reduce) and "keeping number of elements the same" (rearrange). This helps in producing more specific error messages when wrong pattern is passed.
In einsum, when an axis is present in all tensors, operation performs independently for all values of this axis, which is principle einops. tensor pattern only identifies correspondence between axes and element's 's position in tensor, but does not affect arithmetic operation; this ensures that input and output patterns can be changed independently. In particular, rearrange is an arithmetic identity (same value returned), but usage of different input and output patterns turns it into a rather universal tool for changing tensor shape/layout.
Proposed notation addresses different problems of the mainstream approach:
Both input and output are described in the operation definition: tensor dimensionality and expected order of axes. This makes einops-based code more declarative and self-documenting. A user is not required to remember or infer shapes of tensors after every operation.
Input is checked for a number of dimensions and divisibility of corresponding dimensions. The length of dimension checked if provided.
A tensor structure cannot be broken by design, because the notation connects input (or their constituents) to output axes.
Axis enumeration is not used, so no way to make one-off mistake.
Users do not need to compute permutation of axes, those are computed from a pattern.
einops notation alleviates the need to track a tensor layout with patterns.
einops and einsum "document" inputs and outputs, simplifying inference of tensor shapes from the code for other tensors that are not direct input/output of einops, but are interacting or computed using direct inputs/outputs (an example is given in Appendix E).
We show versatility of einops by expressing common numpy (np) operations in Listing
np. transpose(x, [o, 3, 1, 2])
rearrange(x, 'b w c -> rearrange(x, 'h w -> (h w)
w')
np. reehapexxx [x. ssssaee(o=x.ssppellll shsaa][[2]]
np.squeeze(x,
rearrange(x, '() w
ct)
np.expand_
np.expand_dims(x. dims_ddissx -1)
rearrange(x, 'h w c >> c ()') rearrange(Ir, 8, b], 'c w >> h c() rearrange(Ir, g, bl, 'c hhww >> (c h) w¹)
np.stack([r, g, b], axis=2)
ap.concatemate(1r.
b], axis=0)
'Post-factum we can confirm choice: search over Githuh shows that rearrange, the most restrictive on possible patterns, also accounts for the majority of usages.
'Soe examples list inputs, details see in Appendix B.

---

Published as a conference paper at ICLR 2022
np. flatten(x)
rearrange(x, "b >> (b (bt)
np. ax 0,
rearrange(x, 'b -2 b
10 left, right np.split(image, 2, axis=1)
rearrange(x, 'h (Ir w) ->
w c', Ir=2)
11 even, odd = x(:, 0::2], x(:, 1::2]
rearrange(x, 'h (w par) -> par h W c', par=2)
12 np.max(x, [1, 2])
reduce(x, "b
->. ".'max""
mp.mex((x)
reduce(x, "b <>', 'mean')
14 np. np. mean(x, axis=(1, 2), Trepd)
reduce(x, h
()( c', 'mean')
15 np. reshapex,,,,, 1-1, ml)x.mx(xxis==)
reduce(x, '(h 2) -> h', 'max')
16 repeatxx,,, 2, axis=1)
repeat(x, 'h
(w 2)")
np.tile(x, 2, axis=1)
repeat(x, 'h h (2 w) repeat(x,
18 np.tile(x(:, ;, np.newaxis], 3, axis=2)
3¹)
Listing I: Correspondence between numpy and einops operations.
5
CASE STUDIES
We again consider fragments from the vision permutator (Hou et al., 2021). Two examples below only differ in what axis is mixed.
vision permutator mixing along
x_h = xershhppe(B, H, w, N, permutuee 3, 2, 1, r).reshape(B, N, W, H*S)
x_h = proj_h(x_h)) (eexx_hh.reeeheeee reshape(B,, N, W, H, p). 3, 1, 4) reshappe( H, W, c)
vision permutator mixing along
x_ x.reshape(B, w, S). >)ppemuueee,, 3, 2, 4) reshapeee H, W+S) __w proj_w(x_ww.resheeee
H, N, w, s).permute(o, 1, 3, 2, reshhapp)) H, W, c)
While four operations were updated simultaneously, in einops counterpart changes are limited to swapping and before and after projection, because layouts of input and output are disentangled and changes in one do not propagate to the other. And to remove batch dimension, one just removes axis from patterns (this is also generic property einops).
mixing along h x_h = rearrange(x, 'b w (n s) -> b n (h s)', s=S)
# einops:
x_h rearrange(proi__xx.h))
"b (h s) h h (n s)', s=S)
einops:
mixing along We swapped and before and after projection rearrange(x, 'b (n s) -> ->b (w s)'s, s=5)
x_w
x_w = rearrange(proj__(xx.u
'b n (w -
s)', s=5)
The next fragment is derived from OpenAl's implementation of Glow (Kingma & Dhariwal, 2018). As other neural flows, Glow includes multiple rearrangements.
def unsqueezezd(x, factor=2):
assert factor >" if factor return
shape = x.get_ shat__shape()
height int(shape[=])
width int(spae(22)
n_channels int int(shape[3])
assert n_channels >= 4 and n_channels & ==
x = ref.reshape(
(-1, height, width, int(__channels/facton
factor, factor))
* [o, 1, 4, 2, 5, 3])
x tf.reshape(x, (-1, int(height+factor), fattii) into__chameetsffeoorr2))))
return
same in einops, function introduced
rearrange(x, 'b (c h2 w2) (h h2) (w w2) c', h2=factor, w2=factor)
As for the original implementation, function name is confusing and non- descriptive. To reflect the actual transformation, name should be changed to rearrange_be troduce function, arguments describe input, output and the transformation itself. This
Sastee____ssue_____eeeeeeee swweeezzeehannelsseee Laaiumel____n_uussueeeeeeee and einops alleviates necessity to in-
7

---

numpy Published as a conference paper at ICLR 2022
example also demonstrates how einops performs transformations, under the hood it makes the same sequence of transformations (reshape-tramee transpose- (epposeer as in the original code. Reverse rearrangement (also used in Glow, see Appendix F) requires detailed analysis before implementation as no arguments are shared between original and reverse rearrangements. einops a pattern can be reversed by swapping input and output parts of the pattern.
In Appendix E we analyze and rewrite a larger fragment of code: a multi-head attention module derived from the popular implementation (Huang, 2018).
IMPLEMENTATION DETAILS
We implement proposed notation in Python package einops."
Support of multiple frameworks. einops supports a diverse set of DL frameworks (pytorch, tensorflow, chainer, jax, gluon) as well as frameworks for tensor computations: numpy, cupy (Okuta et al., 2017). We refer to them as backends. The major challenge in the support of multiple backends is the absence of common API: even simple operations like repeat, view, transpose defined inconsistently.
np, transpose(x,, [2, 0, 1l)
x. transposse 9, 1)
numpy
tf. transpose(x, [2, 0, 1])
tensorflow
K.permute dimensionssx, [2, 1]) keras
nd. transpose(xx, (2, 1))
mxnet
x.permute(2, 0, 1)
torch
rearrange(x,
"h W - thw'
einops (any backend)
This inconsistency makes projects like einops more valuable for users, as they minimize framework specifies when users do not need it. Proposal (Data-Apis, 2021) may help in convergence shared API and simplify development of cross-framework applications, including e einops.
Backend recognition in einops does not wrapped imports, which commonly used in Python to handle optional dependencies:
def is js_numpy_teemee (x): einops does not use this approach
try:
import numpy as np
return isinstance(x, nd.rrayy
except ImportError
return False
Instead, einops keeps dictionary that maps a tensor type a backend. When the type is not found in dictionary, pass through backends is done, and before importing the module, einops confirms that the backend was previously loaded in sys.modules, since einops will not receive tensors from non-imported modules. The main reasons for this strategy are: a) some backends take a lot of memory and time load (e.g. tensorflow), which may result in an unforeseen usage of resources if backend installed but not used; b) it possible that a backend is installed incorrectly (e.g. wrong binary) and import drives to non-catchable segmentation faults; c) checking backends one-by-one incurs unnecessary overhead, which we skip by dictionary lookup.
einops has shared logic for parsing and conversion into operations for every framework. This conversion is very efficcint: every einops pattern is converted into at most 4 tensor operations by backend, even dimensionality of tensor large.? Overall, einops adds marginal overhead on top of the DL frameworks, Appendix G.
Support for backends. Each backend represented by a proxy class that provides minimal set of operations necessary to manipulate tensors. In addition, several supplementary methods are required and implemented to allow generic testing: most of tests can be run against any backend. In addition,
Einops package is available online at ht sttps:///a /tg::///tthu.... www. hat///rrggzzokkkokoo einops "There unavoidable exceptions: backends cannot reduce multiple axes
8

---

Published as a conference paper at ICLR 2022
einops implements layers for backends with layers support. This codebase heavily relies stateless operations and has some backend-sye backend-specific code.
Caching plays important role in einops high performance (Appendix H). There are two layers of caching: i) cache for a pattern and provided axes; ii) cache for a pattern, provided axes, and input shape. When the same pattern is applied to a new shape, only shape verification and computation of unknown axes sizes is done. When a pattern is applied to input of the same shape (quite typical for iterative algorithms, and very common in DL), einops only executes sequence of commands with cached parameters.
Exceptions detailed to provide information about performed operation including pattern and provided axes sizes to simplify debugging.
DISCUSSION
Speaking of criticism, einops is sometimes described as sstringly-typed We should point that does not make einops less reliable compared to the mainstream approach:
numpy.transpost transpose: [ndarray, Tuple[int]] -> np.nrarray
rearrange: Indarray, string] >> ndarray, (ndarray can be polymorphic)
So input and output types are not any different except for string not being tuple. In frameworks tensor types do not describe number of dimensions, and existing type system cannot set restrictions on shape components. Static analysis identically ignorant to mistakes patterns and e.g. axes order. There is no static check that prevents users from passing tuple with wrong number of components, or repeats, or just ridiculously large numbers. Interestingly, numpy einsum can accept axis indices not string pattern, but it makes operation less expressive and almost never used.
einops does not enforce any alignment of shapes/axes across operations, thus does not provide integrated analysis/tracking of shapes. Unfortunately, there is currently no design that can accomplish this goal either: different flavors of labeled tensors remain prototypes. A successful design should satisfy multiple requirements that are rather challenging and seem to be irresolvable without new language (or new language features). On the other hand, to become practically valuable, the system should not be too sophisticated. If resulting solution is significantly harder integrate compared to running every module with several test inputs, it is unlikely to get wide adoption.
We performed initial exploration of notation applicability in 2018 when einops was open-sourced by reimplementing a set of diverse fragments from popular DL implementations across different tasks and domains (see details in Appendix I). Since its initial release, adoption of einops steadily grows, it is currently used in more than 1000 public Github projects including the contributions from widely known AI laboratories. This important evidence (discussed in Appendix confirms design choices and ultimate suitability of notation for needs of research in deep learning.
einops notation is not tied to a specific implementation or even language: it can be implemented in more efficient languages." 8 Implementation also can use lower-level primitives and optimizations provided e.g. by TVM and MLIR, not framework-provided abstractions. Both directions have a potential to boost einops performance. Alternatively, notation can be introduced to lower-level primitives.
8 CONCLUSION
We introduce einops minimalist notation for tensor manipulations that leverages patterns for describing structure of input and output tensors. einops mitigates a number of issues common to the conventional tensor manipulation routines, represents a number of commonly used functions with small API, and makes code more readable and reliable. We implement notation in Python package that provides identical API across a variety of tensor frameworks. fross-framewor support, expressiveness of code, additional checks and simple integration into existing projects make einops a convenient tool for researchers and engineers.
*e eeioops independently ported to Rust language http:://los..///tinng//a/ttsttta

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
We thank M. Cramner, C. Garcia, E. Glazistova, Grachev, D. Ivanov, J. Henriques, Hoyer, T. Likhomanenko, A. Molchanov, S. Ramakrishnan and others for participating in user studies and/or discussions of notation and API. Multiple improvements the project documentation were made by C. Garcia and other community members.
We thank Phil Wang for publishing numerous high-quality implementations of new architectures using einops and introducing practitioners to the notation by demonstrating its usage in complete projects.
We are grateful to A. Chernyavskiy for feedback on initial draft of this manuscript and to T. Likhoma- nenko for multiple suggestions on the paper, help during review process, and funding initial stages of this project.
We thank anonymous reviewers and program chair for discussion that helped us to improve and strengthen the paper.
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https: htop:://www... www.ttetsooofooooooooooooo org/. Software available from tensorflow.org.
Olah,
Charu C Aggarwal et al. Neural networks and deep learning. Springer, 10:978-3 2018.
Yoshua Bengio, Ian Goodfellow, and Aaron Courville. Deep learning, volume 1. MIT press Massachusetts, USA:, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githin com/ homm/googl/aax. le/jax
Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogencous distributed systems. arXiv preprint arXiv:1512.01274. 2015.
Eric M. Christiansen and Stephan Hoyer.
l.bbetrib.labeled.. labeled_tensor.
https://github http://ttt...comm/eeooosoo//ttrr t/ com/ httr///aalooossaaaffffffoooeeeaa
2016.
[Online; accessed 1-Oct-
2021].
Data-Apis. Consortium for Python Data API Standards, Arry API standard. https://github. 2021. [Online; accessed 1-Oct-2021].
David Foster. Generative deep learning: teaching machines paint, write, compose, and play. O'Reilly Media, 2019.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Marchaat, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585 September 2020. doi: 10.1033//41588-000 -20-222--9-22 2649-2. URL httpss://oll
10t.....111111111111 $41586 020- 2649-2.
10

---

Published as a conference paper at ICLR 2022
Qibin Hou, Zihang Jiang, Li Yuan, Ming-Ming Cheng, Shuicheng Yan, and Jiashi Feng. Vision permu- tator: A permutable mlp-like architecture for visual recognition. arXiv preprint arXiv:2106. 12368, 2021.
Stephan Hoyer and Joe Hamman. xarray: N-D labeled arrays and datasets in Python. Journal of Open Research Software, 5(1), 2017. doi: 10.5334/jors. 148. URL http://dot org/10. 5334/ jors. 148.
148.
Yu-Hsiang Huang. Implementation of "attention is all you need" paper. https: //github. wttp://aaooreee111oo/tttttttooon 2077515
2018. [Online: accessed 1-Oct-2021).
Kenneth E Iverson. A programming language. In Proceedings of the May 1-3, 1962, spring joint computer conference, pp. 345-351, 1962.
John Jumper. Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021.
Diederik P. Kingma and Prafulla Dhariwal. Glow, official obbo:oofficialmmmmeeeeentttoonn https: com/ openai/s glww/ httpw//w///hh/o//o http:///www./.......//////////a. 2018. [Online; accessed 1-Oct-2021).
Olcksii Kuchaiev, Jason Li, Huyen Nguyen, Olcksii Hrinchuk, Ryan Leary, Boris Ginsburg, Samuel Kriman, Stanislav Beliaev, Vitaly Lavrukhin, Jack Cook, et al. Nemo: a toolkit for building ai applications using neural modules. arXiv preprint arXiv: 199.09770 2019.
Hanxiao Liu, Zihang Dai, David So, and Quoc Le. Pay attention mlps. Advances in Neural Information Processing Systems, 34, 2021.
Matlab, Toolbox, symbolic math and others. Mathworks Inc, 1993.
namedtensor. namedtensor python package. htttp:///gittub..cc///avrrrrrr namedtensor, 2019. [Online; accessed 1-Oct-2021).
Numpy. Numpy broadcasting rules. httpss//numpyn broadcasting.u html, 2021. [Online; accessed 1-Oct-2021].
Ryosuke Okuta, Yuya Unno, Daisuke Nishino, Shohei Hido, and Crissman Loomis. Cupy: A numpy- compatible library for nvidia gpu calculations. In Proceedings of Workshop Machine Learning Systems (LearningSys) in The Thirty-first Annual Conference on Neural Information Process- ing Systems (NIPS), 2017. URL htttp://Leaannngsyy. wwwgggnisi//asseeeeeeeeeeeeeeeee p papers/ paper 16. pdf.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc. E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http:::/pppprs...eiiss..co.oooo library. ottpeeeeeeeeeiiipprrrrrrrrreeeee
pdf.
Adam Paszke, Daniel D. Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Mathew J. Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. Getting to the point: Index sets and parallelism-preservvin autodiff for pointful array programming. Proc. ACM Program. Lang., 5 (ICFP), aug 2021. doi: 10.1145/3473593, URL https://doi.org:
11011144/3377 34735933
Pytorch. Documentation for Pytorch named tensors. https: http:://pptooocc.rr///ooo////// named tensor. hamed_tensor. html, 2019. [Online; accessed 1-Oct-2021).
11

---

Published as a conference paper at ICLR 2022
Seiya Tokui, Ryosuke Okuta, Takuya Akiba, Yusuke Nittani, Toru Ogawa, Shunta Saito, Shuji Suzuki, Kota Uenishi, Brian Vogel. and Hiroyuki Yamazaki Vincent. Chainer: A deep learning framework for accelerating the research cycle. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp- 2002-2011, 2019.
Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mlp architecture for vision. Advances in Neural Information Processing Systems, 34, 2021.
Hugo Touvron, Piotr Bojanowski, Mathilde Caron, Matthieu Cord, Alaaeldin El-Nouby, Edouard Grave, Armand Joulin, Gabriel Synnaeve, Jakob Verbeek, and Hervé Jégou, Resmlp: Feedforward networks for image classification with data-efficient training. arXiv preprint 2rXhv:2105.08 2021.
Guido Van Rossum and Fred L. Drake. Python 3 Reference Manual. CreateSpace, Scotts Valley, CA, 2009. ISBN 1441412697
Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. Tensor comprehen- sions: Framework-agnostt high-performance machine learning abstractions. arXiv preprint arXiv:1802.04730, 2018.
Nicolas Weber and Michael Goesele. Matog: Array layout auto-tuning for cuda. ACM Transactions on Architecture and Code Optimization (TACO), 14(3):1-26, 2017.
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions vision transformers, In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 22-31, 2021.
12

---

Published as a conference paper at ICLR 2022
A DISCUSSION OF SUPPORT FOR COMMON DL BLOCKS IN LABELED TENSORS
While most of labeled tensors packages were targeted at deep learning, none of them proposed integration of DL-specific blocks like convolutions. For illustrative purpose let us discuss different possible scenarios how labeled axes could be added to convolution:
Should convolution stick to existing behavior and rely axes order not names? If so, get additional checks, but create potential problems since axes names are required in other operations (e.g. reductions), but users should drop the names before every convolution, and reorder if necessary. Next, users have to set axes for convolution output, in agreement with input axes labels.
Should convolution rely convolution may require input to have dimensions ("channel", "height", "width") while all the other dimensions could be considered "batch-like" and preserved in the output? In this case, almost all the tensors will have the same axes names, thus name checks become meaningless. If one decides to apply attention or e.g. recurrent unit, each switch would require renaming axes.
names? If so, should axes names be fixed? For instance,
Should output names be predefined? Thus all tensors produced by convolutions would have the same labels, still preventing meaningful checks. That seems to produce more potential confusion than help.
Should input axes names propagated to output? If yes - which names should propagated? Output channels at first seem to be disconnected from input ones. However, in practice opposite is equally frequent: in depth-wise convolution each input has one-to-one correspondence with output. Setting up common rules for propagation of "height" and "width" labels meets similar questions and exceptions: compare padding= same and padding= valid'.. vadding="valid".
Similar questions arise for other "building blocks" of DL.
The broader problem is i.c. sets restrictions that help in detecting or preventing coding errors, while introducing limited performance and coding overhead.
up with concept of axis "label" that is indeed helpful in development,
All design choices taken (including support of various DL blocks) should be in agreement with each other, guided by some high-level logic that works for multiple models and applications. This hardly feasible without well-defined concept of axis "label", which was not proposed by implementations.
B EINOPS AND EINSUM
einops notation is largely inspired by numpy. einsum and extends it for single-tensor transformations: in Listing 1 only 2 operations of 17 (transpose and swap axes) can be implemented with einsum. We design interfaces in einops to be aligned with numpy.einsum for simultaneous usage. In the case study of multi-head attention, Appendix E, we demonstrate how they should interact together. Overall, einops implementation deviates from numpy. einsum the following:
einops supports arbitrary reductions (max-, min- sum-, mean-, logaddexp-reduucionss as well as callables to support any other reduction) while einsum can be used to perform sum-reduction only.
While einsum allows only one-character names for axes, einops supports multi-character ones: digits, underscore, and arbitrary capitalization are allowed, while space is used as delimiter. For example, einsum pattern 'hw -> c hw" transposes 3-dimensional tensor, while this an operation 2-dimensional tensors in einops.
While einsum implementations suffer from non-aligned support for spaces between names and capitals, einops notation beckend-independen
In addition to these discrepancies, we introduce several new concepts in einops which are absent in numpy. einsum:
Composition and decomposition of axes, new core functionality described in Section 4.
13

---

Published as a conference paper at ICLR 2022
Unitary axes: in any pattern
to reflect axes of size
Specification of axes size and verification of shapes and divisibility.
Anonymous axes: axes that are present only on Ihs or rhs can be specified with just size in the pattern. For example: repeat(x, 'h -> h c', C=3) can be repeat(x, 'h -> 3
Note that unitary axes are not sub-case of anonymous axes.
Introduction of new axes with einops. repeat by repeating elements along the axis. We demon- strate that this operation allows einops to support repeats and tiles.
List imputs. einops accepts a list of tensors with identical shapes and dtypes, which are channel'). )) This
stacked. For example, rearrange([r, 8, alphal, "channel extends e einops to support tensor stacking and concatenations.
Layers. einops functions have layer counterparts, e.g. einops. layers. torch. Rearrange( c -> (c h w) can be used inside Sequential as one of the modules." ?
and at the same time we do not implement some features of sumpy.einsum:
Axes repeats on lhs. While einsum allows writing e.g. 'ii->" to get matrix trace dropped this einops to prevent coding mistakes.
Multi-tensor transformations. While there is an einops layer that operates on two tensors, the core functionality of einops used by most users is one-tensor transformations. However, einops novelties (new axes, anonymous axes, unitary axes, specification of axes sizes) can readily be applied extend einsum.
Implicit indexing. This feature relies on characters sorting, does not encourage descriptive axes names as those define output axes order, and makes output implicit, not explicit. Implicit indexing also cannot be applied to einops. reduce and einops repeat. There is negligible usage (< 1%) of list-inputs syntax numpy. einsum, and is not added einops.
C FORMAL GRAMMAR OF EINOPS PATTERNS
To accommodate differences, covered in appendix B, but still keep similarity with numpy. einsum, einops patterns follow a formal grammar defined in Listing 2. This grammar significantly evolved from the prototype stage of einops notation, when we performed user study described in Ap- pendix D.
pattern = tensor_ censor_pattern, '->', mensor_patter... tensor pptttrr_ [delimiter], axis_group axis "(", [delimiter], ")" I [delimiter], axis, {delimiter, axis), [delimiter]
transformation fortttaioonnnnteeeen
group, ('', delimiter, axis_grouph, [delimiter]
axiss '...' unitary_axis anonymous axes unitary_axis '1'
_eme_axxs___
de delimiter
named_axis axis <any python identifier that does not start or end with underscore> anonymous_axis <natural number greater than one> ;
Listing 2: Extended Backus-Naur Form grammar of einops pattern notation for a single tensor.
D CONFIRMATION OF READABILITY AND INTUITIVENESS
At the time of the first einops prototype, the idea of using operations with named axes tied to operation itself not to a tensor was not around within DI.. Besides, among DL frameworks only pytorch had reasonable support for einsum (with multiple performance issues reported).
That is why at the prototype stage we developed test to probe several design choices and to answer following questions:
* Is the notation easy to pick up?
In examples from pytorch documentation (otherwise Sequential) modules are implemented using custom "forward" accomplish this transformation with x.view flatten.
14

---

Published as a conference paper at ICLR 2022
Does the notation require any specific introduction?
Can users recognize operations without context?
Specifically, 8 subjects with sufficient (>l year) exposure tensor programming in least of the DL frameworks were asked to complete the following questionnaire: subjects were not provided with any description of function, implementation, or other context."
All imports are ignored. What is the output of each print statement? torch.zeros(1o, 20, 30, 40)
transpose(x, 'bhwe->bchw")
print(y.shape) Q1
transpose(x, 'bhwe->bc(hw)")
print(y.shape) Q2
y transpose(x, h1, "bidCc.,,,w))-s0hhhh///... hi=2, w1=2)
print(y.shape) # Q3
transpose(x, 'b(h, bochhhht/w...m)ee-hhhhhhnces h1=2, w1=2)
printty. Q4
transpose(x, 'bhw(c,g)->gblwc, g=2)
printyya.shape, y2.shape) Q5
transpose(x, 'bishat->bib2st")
print(y.shape) # Q6
operator is matrix multiplication
transpose(x, "bchw=>(bbw)e") rarch.rand(20,50)
print(t.shape) Q7
transpose(t, "(bhw)c2->bczhw', www.xx.shapp print(y. Q8
first operand is
transpose(t, "(bhw)c=->bc2hu' b,3e, h=1e)
print(y. # Q9
This user study had some additional ideas that were rejected later: comma as delimiter within parenthesis (non pareethhhhed commas conflict with arguments (e.g. b b_hw=x.shapp Even some participants who correctly answered corresponding questions (5 participants of 8 in both cases) reported they had initial confusion with these features.
and parsing of multiple shape
Only two participants were aware about einsum function but not used that in their coding practice. Questions (QI, Q2, Q6, Q7, Q9) were correctly answered by all participants. This confirms that notation based on named axes intuitive and can be picked up without a special introduction. Post-questionnaire interview showed that participants could describe a particular context when an operation (e.g. Q1, Q3) could be used.
While single-character variables names (with optional digit) did not cause confusion, have replaced this due to conflicts with potential future features (c.g. unitary axes). Function name transpose was also changed to rearrange avoid conflicts with existing operations in different DL. frameworks.
However the real value of design choices in einops can be confirmed in long time-scale and within large codebases, settings that can't be captured by small user studies.
Some participants had no experience with pytorch and they were informed that x shape [1,,20,39, 30,,30,488) 48].
15

---

Published as a conference paper at ICLR 2022
E CASE STUDY: MULTI-HEAD ATTENTION
Below we provide shortened implementation of multi-head attention based on Huang (2018). For brevity, we remove weights initialization and mask support.
class Gellebootrrodcccctteeiion...mm
def __.init__ temperature, attn attn_dropout=o.11 superr))._.iiit__.)
self. telpratrerr temperature
self.dropout nn. (attoutttttttnttootoa self.softmax = nn. (oiiimmmddd))
def forward(self, 9, v):
attn torch. bmm(q, k. transpops, 2)) / self. temperature attn self.softmax(attn)
10
11
attn self.dropout(attn)
output torch. tattt,, v)
14
return output, attn
15
16
class Mulcilecadkteentiionm...n Modulle)
18
__.init__sseff, h_head, d_model, d_k, d_v, dropout=e.1): super()...init__] self.n_head n_head self.d_k d_k self.d_v d_v
19
20
23
24
self.w w_qs on.Linear(d_model, n_head d_k)
25
self.w_ko wskk Liinear(d_modee n_head * d_k)
self.w_vs nn.Linear(d_modde n_head d_v)
27
self.attention self. Leyerlayer_nor _orm nn. LaayerNormmd_moode
e.5))
28
29
self.fc nn.Linear(n_her head d_v, d_model)
delf.dropout nn. Dropout(dropoutt
32
33
def forward(self, 9, k, v):
34
d_k, d_v, n_head self.d_k, self.d_v, nelf.n_head
sz_b, len_q, q.size() sz_b, len_k, k.size() sz_b, len_v, v.size()
37
38
residual
41
42
q = self. w_
len_q, n_head, d_k)
43
k = self. (kks wv..ks(k).vie((s2z view(sz_b, len_k, neadd d_k)
self.w_vs( view(sz_b, len_v, n_head, d_v)
46
9 9- permutee, 0, 1, 3). c).ccotiguouu(().vr -1, len_q, d_k) # (n=b) x lq x dk
47
k = permute,, 0, 1, s).contiguous(( ())tuoous().iie) -1, len_k, d_k) # (n*b) x 1k x dk ().contiguous().vieee)) len_v, d_v) # (n*b) x lv x dv
48
v v.permute(2,
50
output,
self self.attention(q, k,
51
52
output output. seeutt.veewnn_hee n_aedd sz_b, len_q, d_v) output output. permute(,, 2, co.contiguous( view(sz_b, len_q, -1) # b x la (n*dv)
54
55
output self. output self.layer Lartr_yyoeomomuppuu residual)
56
output, attn
16

---

Published as a conference paper at ICLR 2022
Original implementation demonstrates number of issues that previously discussed: unchecked and hard-to-track transformations, necessity to keep comments about the shape. Frequent usage of -1 in reshapes makes it simple to introduce errors. The critical part of computations (attention) is offloaded to a separate module, which does not check input and which not documented: dimen- sions of inputs and outputs not defined, just the connection between them. Thus, provided weleeDDoorrduucttttee module.
cannot be considered as an independent, self-containing or reusable
Inlining of Cealeddooooouiccctteeetoooo tAttention inside Mult MeltieedAtttettinn could improve the situation, but would also make the problem with shapes more obvious, as more comments would be required for inlined variables.
We can compare that with einops implementation, where all computations are done in single module. Each axis is easy to track throughout the code. An axis index is needed three times (explicitly in sof tmax, implicitly in and layer_norm) but those are easy find from the code without any additional comments which demonstrates how einops "implicitly annotates" code in the proximity.
class Multihheadkteentioneeell inrreunnn..Moooooo Module):
def ___iiit__seell n_head, d_model, d_k, d_v, dropout=o.1):
self.n_head n_head
self nn.Linear(d, _a__eerrdd_mdel,,, n_head d_k)
self.w_ke an.Linear(( (d_oded,_ n_head d_k)
self.fc (alLineardd_moeee _mmoeel,, n_head * d_v)
self.w_vs mn.Linear(d_
d_v, d_ m_model)
11
self.dropon (n.proooot(psrrooou)) self.attn_dropout an..Droouttppe.. self.tayer_norm nn. ww..ayyrrorrmm__mde
12
13
def forward(self, 9, k, v):
16
residual
17
q rearrange(self..wogg wsteesssellwwwas()..)))))_))))() "b (h k) >> b 1 k', haself.n_heal
18
k = rearrange(self.. ww.ksssllllwwwkkkk_k_ "b (h k) >> h k', beself.n_headd
rearrange(seelf.wwssV t
v', h=self.n_head)
20
attn torch.einsum( "bbbk,,btt [q, k]) / np.ssrrt(q.shapee[1)
21
attn self.attn, wtt.attt.crppottttttttsstttaaaaa
22
output torch.einsum("hblee hblthuutttttttt [attn, v])
output rearrange(outpput blbvv b (h v)
output self. gelF.ddooottsstf.fccoooooo output self. Lawee_rrrroomoooteetto residual)
25
26
return output, attn
17

---

Published as a conference paper at ICLR 2022
F SQUEEZE EXAMPLE FROM GLOW WITH REVERSE TRANSFORMATION
Squeeze and unsqueeze implementations are derived from OpenAl's implementation of Glow (Kingma & Dhariwal, 2018).
def squeezezd(x, factor=2):
factor factor 1:
return x shape x.get_shape()
height int(shape[11)
width int(shape(z])
n_channels int(shape(3))
assert height % factor == and width & factor ==
x tf.reshape(x, [-1, height//factor, factor, width//factor, factor, n_channels])
x tf..transposexx, [o, 1, 3, 5, 2, 41)
x = rf.reshapexx, [-1, height//factor, width // factor, o_chameels-fcctorrfactoo
return
def unsqueezezd(x, factor=2):
assert factor >= if factor return x
shape = x.get_shape(
height int(shape[1])
width int int(shae((2)
n_channels int(shape(31)
assert n_channels > 4 and n_channels &
x = tf. tf.reshape(
(-1, height, width, intn_chanmels/factore/a factor, factor))
transranpose [o, 1, 4, 2, 5, 3))
x tf.reshape(x, (-1, int(tneigtt-factor) int(witth in((__hhamells/fccoor+))))
return
einops counterparts for both functions can be written as
rearrange(x, "b (c h2 w2) -> (h h2) (w wz)', h2=factor, w2=factor)
rearrange(x, "b (h h2) (w w2) >> (c h2 w2) h w', h2=factor, w2=factor)
where one can readily see from the patterns that the second operation is inverse to the first one.
G EINOPS PERFORMANCE
To demonstrate that the overhead brought by einops on top of DL framework is negligible we measure performance of several case studies. We compare original pytorch implementations and their einops versions in the following scenarios, see Table 1: CPU or CUDA backends, with enabled or disabled JIT (just-in-time compilation), different input tensor sizes. In our performance benchmark we use einops 0.3.2, pytorch 1.7. 1cu. 10, CUDA 11.0. We use AWS EC2 p3 p3.2xlarge instance for benchmarks. The average time of the forward pass is measured in milliseconds by IPython's module timeit. For multi-head attention case study, see Appendix F,, einops implementation uses pytorch einsum operation which gives additional overhead. That why, for this case study we consider einops implementation with and without einsum. For the of unsqueeze, we popular open-source port to pytorch https: / /gggttth....... cww homm/hhaaujinn// glww- pytorch.
From these benchmarks even e einops-rich code works a similar range of speeds pytorch-only. In particular, unsqueeze2d consists only of einops operation, shows similar speed both on CPU for various input sizes and GPU with larger input sizes. Since computationally most expensive operations are convolutions and tensor products (linear layers, attention). not rearrangements. this overhead will typically have marginal contribution the result. Thus, attention and permutator closer to
18

---

Published as a conference paper at ICLR 2022
practical use cases, and demonstrate that difference in performance becomes notable only for small inputs when GPU is used.
Table 1: Performance comparison between original pytorch implementation and its einops version for different case studies. Performance is measured ms via IPython's 's timeit module.
CPU Backend
CUDA Backend
Case Study
Input Size
Impl.
w/o JIT
JIT
w/o JIT
JIT
cinops
21.85=0.04
21.13+0.03
1.30=0.00
1.05=0.00
(32. 64, 512)
w/o einsum
21.31#0.04
20.9920.06
1.1420.00
0.9040.00
original
21.09=0.03
20.85=0.05
1.02=0.00
0.71+0.00
cinops
48.0620.06
47.37:0.05
1.39=0.00
1.30+0.00
(32, 128, 512)
w/o einsum
46,94:0.04 04
46,28:0.05
1 1.28=0.00
1 1.29:0.00
attention
original 46.87=0.08
46.33=0.04
1.29=0.00
1,28+0.00
einops
150.40=0018
149.53+0.17
2.78+0.00
2.79+0.00
(32, 256, 512)
einsum
146.01+0.20
144.9620.19
2.75+0.00
2.7640.00
original
145.19=2.07
14.294=0.21
2.75=0.00
2.75+0.00
cinops
495.31#1.19
993.90=0.60
6.51#0.00
6.53+0.01
(32, 512, 512)
w/o einsum
488.37=0.93
187.1620.96
6.47=0.00
6.48+0.01
original 491.20+2.79
488.6440.23
6.45#0.01
6,4620.00
einops
7.05=0.01
7.03=0.01
0.52=0.00
0.52+0.00
(32, 32, 32, 32)
original
7.0820.02
7.05=0.02
0.44
0.40+0.00
permutator
einops
237.53+1.67
237.05=0.24
2.51#0.00
2.56+0.00
(64, 64, 64, 64)
original
234.57=0.17
234.66#0.21
2.50=0.01
2.54+0.00
einops
1.6840.00
1.73+0.00
0.07=0.00
0.13+0.00
(32, 32, 32, 32)
original 1.66+0.00
1.72=0.00
0.05=0000
0.09+0.00
unsqueeze2d
cinops
16.91=0.02
17.01=0.04
0.14=0.00
0.1000.00
(32, 64, 64, 64)
original 16.89=0.03
17.14=0.02
0.12=0.00
0.14
cinops 129.53 0.05
129.95=0.
0.71
0.76+0.00
(32, 128, 128, 128)
original 129.16#0.16
131.44#0.12
0.69=0.00
0.72+0.00
H ROLE OF CACHING
To estimate the role of caching we conduct study with very small numpy arrays, where we generate 10d patterns performing the same transformation using IPython's timeit functionality. Cache can't accommodate all these patterns and runs syntax parsing, validation and inference of dimensions.
from einops import rearrange import numpy as np
patterns
f'i(i} (j k) 1 >> i(i} (k 1) (m ...)'
for i in range(10_oog
x ne.os [2, 2, 2]) # small tensor
caching mechanism is used
xstimeit
for pattern in patterns:
rearrange(x, pattern, j=2)
caching in use
extimeit
pattern = patterns(o] for _pattern in patterns: rearrange(x, pattern, j=2)
19

---

Published as a conference paper at ICLR 2022
Outputs for the two cases show more than 10-fold difference:
without caching: 663 ms + 20.9 per loop (mean std. dev. of runs)
with caching: 37.1 ms + 165 μs per loop (mean std. dev. of 7 runs)
Times were measured on Macbook Pro 2018 using Python 3.9.0, numpy 1.20.3 and einops 0.3.2.
1 EINOPS FLEXIBILITY AND APPLICABILITY
To confirm wide applicability of einops, search for Github code in pytorch that actively relies heww//essapperraaspss///eee
We restrictively consider only popular (> 1>1000 stars) repositories. also included (all three repositories pass the
Official documentation, examples and torchvision previous criterion). We sample 16 fragments that cover a diverse set of applications and architectures and also smaller fragments from the same repositories, Table 2.
Table 2:
Comparison between original implementation and rewritten version with einops measured in number of code lines and characters.
# Lines
# Characters
Fragment
Original
einops
Original
einops
LeNet-like network
19
15
626
352
Super-resolution
17
11
678
478
Style transfer
192
107
LSTM and language modeling
15
16
745
753
LSTM token-by-token
22
1266
990
ShuffleNet
126
43
4007
1932
ResNet
58
42
2106
1658
FastText
24
723
281
CNN for text classification
43
16
1767
799
Tacotron
36
23
818
662
Transformer's attention
78
32
2637
1538
Self-attention GANs
34
19
1506
959
Time-sequence prediction
27
25
1125
1015
Spacial transformer network
36
1099
1012
Highway convolutions GLOW
6
6
256
239
30
5
917
250
These examples rewritten in einops and torch. einsum: interface of models is kept identical to the original implementation to demonstrate that einops does not demand global code adjustments. In addition, introduction of einops layers allows avoid new classes and reuse iequentttaal in four cases (LeNet, Super-resolution, Fast ResNet).
Resulting code in einops is shorter, Table 2, while multiple operations in a single line. Out of 16 fragments, in one case length (measured in number of code lines in number of characters results identical) is higher than in original (LSTM and language modeling), in one case is identical (Highway convolutions), and in all other cases is tangibly lower."
fewer method chains and do not pack
"Sideeto- comparison hetween original code and einops version with comments can be found online http:://wimnssnnkkkkxxaacxxaaaap
20

---

Published as a conference paper at ICLR 2022
J EINOPS ADOPTION
When it comes to estimating suitability of notation (or any other tool) to help in research, one can try to dissect this question into multiple small questions: how many different cases covered by notation? can it be implemented? can it be integrated with existing tools? is resulting code shorter? how fast is resulting code? is modification with notation easier or faster? While we cover in our studies multiple such factors, this "dissection" does not provide holistic view on role and integration into existing workflow and tooling.
einops has taken "test by time" approach. Developed, implemented and open-sourced in 2018 88, got adoption public projects in majority of largest Al laboratories, which a strong indicator of novelty, usefulness and reliability. For context, pytorch claims to be user- iser-focused Paszke et al. (2019). In keynote talk they confirm poor relevance of micro-benchmarks and rely on adoption as a (lagging) confirmation of design choices.
As of November, 2021, Github reports more than 1000 usages of einops by other repositories with increasing frequeney (last 100 were published in 18 days). In the majority of cases they implement approaches or DL models that were proposed after release of einops. Implementations use different DL backends and span different modalities (even non-traditional applications like fluid dynamics).
This inductive validation (if it was applicable to a large number new cases, should work in more) is extremely time-inefficient, but provides the most reliable proof. einops is used in open source projects from AI laboratories, who use and maintain competing deep learning stacks. This forms another important (though less statistically reliable) indication for adoption.
Qutuliine is availahle http://aaamnnnaaa/smmms/22/2/2ee
21


Reviews:
[{"review_id": "vZo-nJ8Ln7h", "reviewer": "Reviewer_gRMH", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "There are numerous issues with the paper, which I will try to summarize below. \n\nI should mention that, since no code was included in the submission, my assessment of this work is based entirely on the paper and the ideas expressed therein. I am not judging the work based on technical aspects related to the library except for those described in the paper.\nHowever, a complete assessment of the paper should also include a code review (see my first comment in the \"Weaknesses\" section).\n\n**Strenghts**\n\n1. I am a great fan of einsum notation for expressing tensor operations and I agree with the authors on the value of extending the notation as they did. I could see myself, as a practitioner, using the \"rearrange\" function (I thought that the example of reshaping an array of images into a 4x4 grid on page 6 was nice).\nAs an engineering feature, the einops package could be useful.\n\n2. Using the library seems to make the code less verbose in the examples provided in the paper. This gives a reasonable advantage in terms of the readability and maintainability of the code (although this is a purely subjective assessment, as I argue below).\n\n3. The fact that the library supports multiple backends is a plus, although I have some concerns about some design choices discussed in the paper (see below).\n\n**Weaknesses**\n\nUnfortunately, I think that there are considerable flaws in the paper, the motivation behind it, and the described implementation of the library. \nI will try to summarize everything in the following points.\n\n1. While I agree that much of machine learning is based on tensor manipulation, this paper is not suitable for publication at ICLR. \nEven if we were to focus only on tensor manipulation as a machine learning exclusive, better venues exist for this kind of paper, like the JMLR Open Source Software track. There, the paper could be reviewed along with the code (which was not included in this submission even if it is a crucial part of the authors' contributions).\nI realize that libraries like TensorFlow and PyTorch had a paper published at some top conferences like ICLR, but it's clear that their novelty and impact on the community was incomparably larger at the time of publication. \n\n2. The motivation for the paper is weak. I will give some examples of claims made by the authors to motivate einops that I don't think hold up:\n    - Users have to remember axes order: this is largely solved by einsum, not einops exclusively.\n    - Control over data layout: einops does not solve the issue, it just gives a simpler interface to control it. \n    - Einops is declarative and self-documenting: einsum too, einops simply extends the paradigm. \n\n3. The contributions of the paper are not enough to consider a publication. \n\n    The paper should describe the notation in detail, but it doesn't. Instead, it is halfway between a description of the notation and a documentation for the package. The result is that it works badly as both. Some examples: \n\n    - The notation is never formally described, for instance as a grammar, and one must look at examples to infer the form of a correct string.\n    - Many features like anonymous axes and ellipsis indexing are essentially ignored, even if they are a key part of the notation.\n    - A comparison with einsum is missing, which confuses the reader and does not highlight some features that are present in einsum but not einops (for example, computing the trace and the implicit notation).\n\n    Also, the authors claim that they \"align the interface with einsum to allow smooth simultaneous usage\" but this is misleading. Einops is clearly inspired from einsum and simply adds some features (while also removing others, as I said above).\nIn this regard, throughout the paper the authors seem to imply that einops solves problems (like that of rearranging) that were previously unsolvable (e.g., second paragraph of page 6). This is also misleading since einops is merely an interface/API to existing functions. \n\n4. The example on page 4 concerning convolution is unclear.\nFor example: when convolutions rely on axes order, the user is not \"expected\" at any time to use named axes for other operations.\nIt is also not clear what kind of \"name checks\" would be prevented by assigning fixed names to the axes in convolution, and how einops overcomes this issue. \nI might have misunderstood this whole paragraph, so please correct me if I'm wrong.\n\n    In this regard, if the interaction with \"neural layers\" is a primary motivation for developing the library, the feature should be discussed more in depth (it is only briefly mentioned once, towards the end, with no explanation).\n\n5. There are many non-rigorous/subjective claims. I have found at least five: \n\n    - Footnote 4: \"New users frequently continue trying to imagine the tensors layout in memory.\" -- On what grounds are the authors making this claim without any references? Have the authors conducted user studies? \n    - Page 9: \"Caching plays an important role in einops high performance\" -- Since there are no performance benchmarks, how is performance quantified here? I assume that the performance of einops is equal to or worse than the backend, given the extra overhead. Is it faster than equivalent einsum operations? Can the authors quantify the \"important role\" of caching?\n    - Page 9: \"einops was also referred once as a good intermediate solution\" -- By whom? How is this whole reported conversation meaningful to the reader?\n    - Page 9: \"Design of such system is much harder than it sounds: previous ideas failed\" -- What does it mean that designing is hard and how can the authors be sure that it wasn't due to their own limitations? Whose previous ideas have failed and in what context? \nI am obviously not implying anything about the programming and engineering abilities of the authors, I am merely pointing out that claims like this should not be made in a scientific paper.\n    - Page 9: \"We observed that einops notation gets picked up for describing tensors with packed dimensions\" -- Again, did the authors perform a user study to make this claim? And, in any case, did einops notation get picked up more easily than einsum notation? If so, why?\n\n**Other minor issues**\n\n6. Figure 1: the example assumes that the reshape is done by hardcoding the dimensions and that the programmer will make a typo that \"luckily\" does not crash. First, we don't know how often this kind of bug happens and, second, this issue can be easily bypassed by accessing the actual shapes (like the authors do on page 6, line 2-3 of the code block at the bottom). I wouldn't use this as a primary motivation for the rearranging function.\n\n7. Page 4, third bullet point: the authors imply that the code should crash, but in fact it is a perfectly well-formed expression that should not crash. Again, the authors should not assume the likelihood of a bug or what the users want to do.\n\n8. Page 7: \"Einops alleviates necessity to introduce a function, as arguments describe input, output, and the transformation itself\". This is a claim in support of the readability of einops, although it holds for einsum already. While I partially agree, einsum/ops notation could also be seen as less beginner-friendly because it requires the users to know the syntax. For example `arr.flatten()` seems (subjectively) more descriptive than `rearrange(\"ab->()\", arr)`, especially if one doesn't already know einsum notation.\n\n**Comments on the implementation details**\n\n9. The choice of automatically detecting the backend from the input is arguable. \nQuoting from the famous Zen of Python by Tim Peters: \"Explicit is better than implicit\" and \"In the face of ambiguity, refuse the temptation to guess.\"\n\n    For example, a situation could happen in which a Numpy array is given as input to a TensorFlow-based user-defined function. \nSince TensorFlow knows to automatically convert arrays to tensors at runtime, this is standard and expected behaviour. \nReplacing the initial call of this imaginary function with an einops-based one would result in the first operation happening in Numpy and not TensorFlow. \n\n    I advise the authors to adopt a paradigm like that of multi-backend projects like Keras in which the users decide explicitly what backend to use (through a config file or by making the backend explicit at import time -- e.g., `from einops.torch import ...`). \n\n10. Can the authors provide an example of how they improved the exception message for broadcasting w.r.t. the one from Numpy? How is it easier to understand?\n\n**Suggestions**\n\nAfter revising the paper, the authors should consider submitting it to an appropriate venue for open source contributions. \n\nThe revision should make it more clear that einops is an extension of einsum and should highlight the concrete additions and limitations of the proposed einops. \n\nThe authors should make the paper more rigorous: the notation should be formalized, the claims about usability should be removed unless backed up by user studies, and the motivation of the paper/library should be rethought. \nThe paper could also focus less on examples (which are more suitable for the documentation).", "summary_of_the_paper": "This paper extends the Einstein summation notation (einsum) of Numpy to introduce some additional features: \n\n- naming axes arbitrarily (i.e., with more than a single character);\n- controlling how groups of axes are flattened (called \"rearranging\" in the paper);\n- repeating axes;\n- reducing axes with different operations other than the sum.\n\nThe notation is implemented in a Python library called \"einops\" that works with different backends for tensor computation (PyTorch, TensorFlow, Jax, etc.). ", "summary_of_the_review": "**Positive aspects**\n\n- The proposed extension of einsum is useful and I can see it having large adoption in the community of scientific computing.\n\n**Negative aspects**\n\n- I don't believe that ICLR is the correct venue for a paper like this. \n\n- The novelty and quality of the paper are rather limited. \n\n- The motivation for the paper is weak and much of the advantages brought by the proposed \"einops\" were already brought by the original \"einsum\" notation. \n\n- The presentation is, at times, confusing.\n\n- There are many non-rigorous claims that cannot be tolerated in a scientific publication.\n\n- There are some design choices that could lead to unexpected problems when using the library. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VlDHW7rSJ2U", "reviewer": "Reviewer_ntk4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n-\tThe proposed toolbox allows for simple and easy to read code tensor manipulation operations and operates with most popular Python frameworks.\n-\tThe usage of indices names is more flexible compared to already existing ones in the einsum operation in numpy and Pytorch, which are known to have some limitations, e.g. number of dimensions is limited by the number of letters.\n-\tIt is very useful for using Tensor Networks where number of interacting tensors are usually big and can have a large number of dimensions. \n\nWeaknesses:\n-\tThe paper does not provide any advance in theory or new algorithm for machine learning. It is limited to introduce a useful and appealing new coding tool.\n-\tThe paper does not mention its application for computing and manipulating Tensor Networks, missing a very important usage for which there is a growing audience eager to have such convenient tool.\n-\tEINOPS does not consider operations involving two or more tensors\n-\tA comparison in terms of computation cost is missing in the paper\n", "summary_of_the_paper": "In this paper, a new Python toolbox called EINOPS (Einstein Operations) is introduced. The proposed toolbox allows for applying tensor operations on a single tensor data, such as reshaping, reducing (max, mean, etc), repeating, permuting, and others; in such a way that it improves code readability and flexibility. The paper illustrates about current issues on available solutions, typically in numpy and shows how EINOPS solves those issues. Also, the toolbox supports widely used frameworks such as Pytorch, TensorFlow and others providing framework-independent minimalist API for tensor structure manipulations.", "summary_of_the_review": "I found the tool very useful for any application involving tensors with Python. The paper is focused on deep learning applications. However, I am sure it is also very useful for dealing with Tensor Networks (TNs), where many core tensors are interconnected. The paper lack of a description of its application to TN contractions, for example. \n\nIt seems that current implementation only consider operations on single tensors. I would suggest to include also operations on two or more tensors, similarly to einsum in numpy and Pytorch, but using a more flexible indices labelling.\n\nA comparison with previously available solutions, via numpy or Pytorch, for example, in terms of computation cost is missing in the paper, making it difficult to evaluate if there is some price to pay.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "V8R-ue0QNU", "reviewer": "Reviewer_p5Km", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First, this paper studies efficient programming paradigms for multi-dimensional array operation, which is an important intermediate layer for modern ML systems, especially for deep learning systems. Personally I tend to buy these advantages of einops claimed by the author, including:\n\n+ providing semantic check for tensor operations; \n+ including high expressiveness and flexibility for tensor manipulating interfaces;\n+ supporting multiple backend runtimes efficiently; \n+ making tensor manipulating code more readable and reliable, and etc..\n\nThese claims sound very promising and valuable as a ML toolkit. \n\nHowever, there are some fundamental concerns I have for the paper: \n\n+ The writing is problematic as an academic paper. Comprehensively, this paper reads like a technical blog, which tries to introduce and advertise a Python library; some statements in this paper are casual, i.e., in the end of Section 7, the author says \"We intentionally omit discussion of user conveniences provided by einops package: anonymous axes, ellipsis, list inputs and neural layers.\", although I would not doubt about it if reasonable amount of evidence (e.g., empirical study or analysis) was presented, this sentence sounds a little ungrounded. Further, the organization of the paper can be polished as well. To be specific, I do not get why Section 2 & Section 3 are split, it seems that both sections are talking about the issues and limitations of STOA systems. \n+ There is a lack of reasonable amount of empirical study to justify the statements about the contribution of the work. For example, the author claims that einops \"significantly improves code readability and flexibility\" in the abstract, I would expect some user-study (which could invite a group of participants to accomplish some programming tasks about tenor manipulation with and without einops and measure some objective or subjective metrics for evaluation) to justify this claim---note that this is pretty standard in first-tier PL/SE volumes, it would be easy to find hundreds of such papers and adopt such methodology for evaluation. The current Section 5 can be considered as good motivating examples, but it is not sufficient to support this claim following the principle of scientific study. The efficiency evaluation of the implementation is also missing, e.g., it is important to learn how much overhead is introduced by imposing einops over the backend systems, e.g., what is the runtime gap between einops implementations and direct usage of the backend systems' interfaces.\n+ Lastly, I feel lost about the discussion about the difference between einops and numpy einsum. As far as I learn, the interface looks similar (at least from coding examples). As so, I would expect some clear and concrete statements about the distinguish components of einops (different from einsum).", "summary_of_the_paper": "This paper introduces einops, a tensor manipulating library in Python, to efficiently support multidimensional tensor manipulations widely adopted in deep learning. \nHighlighted advantages of einops include:  \n+ providing semantic check for tensor operations; \n+ including high expressiveness and flexibility for tensor manipulating interfaces;\n+ supporting multiple backend runtimes efficiently; \n+ making tensor manipulating code more readable and reliable, and etc..", "summary_of_the_review": "This paper introduces an interesting tensor manipulating library in Python, supporting Einstein-notation style operations over multi-dimensional arrays for deep learning. The library manages to run the convenient interface over different popular deep learning frameworks. \n\nHowever, there is a lack of solid empirical study to validate the effectiveness and efficiency of the design and clear discussion about the difference from the existing tool, i.e., einsum. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7o7TpIlF1qZ", "reviewer": "Reviewer_c6LK", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strength of the work.** First of all, the paper provides a handful of case studies with insightful analysis on\nuser experience on the existing design of operators, and then based on those studies it derives a series\nof notations for highly user friendly operator design. By providing various comparison between\nexisting ones on real-world, the paper leads to a convincing proof of the strength of\nthe proposed notation for deep learning researchers' daily usecases, in terms of expressiveness,\nelegance, usability and debuggability.\n\nSecond, the paper comes with solid engineering efforts that offload the computation to underlying\nlibraries, including NumPy, TensorFlow, PyTorch, MXNet, Keras etc. Effectively, with the solid engineering\neffort, the library serves as a meta framework that can be reused across backends with close-to-zero\nengineering overhead. The paper also explores some techniques like caching to reduce the execution overload\nof runtime dispatching.\n\nThird, by the \"stringly-typed\" abstraction of operators, the proposed notation enforces more explicit\nprogramming of the semantics of tensor axes, as well as more runtime checking to catch correctness issues.\nThis is particularly helpful among deep learning practitioners who are already exposed to a set of conventions\n(e.g. \"NCHW\") to write the right code. The typed operators, instead of typed tensors (e.g. NamedTensor),\nprovides the guarantee that users only need to type important part of their code as a drop-in replacement of the existing\ncode fragments without having to worry about refactoring the entire codebase. This is again particularly helpful\nfor deep learning practitioners\n\n**Weakness.** With the proposed notation, compiler-based approaches, e.g. TVM, MLIR, haven't been explored to further boost performance,\nalthough such opportunities are real and tangible. For example in TVM, the Tensor Expression (TE) can be considered as \na generalized form of EinOps, and it could bring extra performance gain to the users. Therefore, it would be desirable to integrate\nwith those compiler backends to assist with potential performance.\n\nThe normal definition of the EinOps notation, while demonstrated in many examples, are not carried out formally.\nFor example, brackets `()` in the notation may serve different purposes under `rearange`, `reduce` and `repeat`,\non the left- and right-hand side of the notation. It may refer to tiling one dimension, or composition of new dimensions,\nor removing a dimension, etc. Another example that needs more clarity is stack/concatenate where there are multiple\ntensors involved, and might be desirable to state formally on the constraints implied in this particular case.\n\nLast, while the notation is very flexible in expressing layout-related operations, it is yet under-explored on\nexpressing the common operators that deep learning researchers may care about, namely, many variants of convolution.\n\n**Correctness.** The reviewer is not aware of any correctness issue in all demos and figures in this paper.\n\n**Clarity.** The paper is overall rich in detailed and insightful explanation. There is some minor clarity issues as pointed out in the weakness section.\n\n", "summary_of_the_paper": "Operators are key elements in deep learning (TensorFlow, PyTorch, etc) or scientific computing (NumPy) frameworks.\nThis paper introduces a generic, convenient and elegant way of massively describing large portion of frequently used operators\nunder a unified framework, EinOps, inspired by einsum, or the Einstein summation convention.\nUnder this convention, key components like multi-head attention in modern neural networks can be expressed\nwith fewer lines of code. The major contribution of this work includes:\n- Presents a set of novel and highly expressive notations to represent hundreds of operators\nin modern deep learning frameworks;\n- Develops a library that interprets the proposed notations and offloads them to existing implementations\nin NumPy or TensorFlow, PyTorch;\n- Demonstrates significantly improved readability and debuggability with the proposed notations\n", "summary_of_the_review": "In summary, the paper demonstrates clearly an elegant and novel notation to represent a large portion of operators in deep learning workloads, and developed a meta-framework with solid engineering. The reviewer recommends to consider the acceptance of the paper as a novel good deep learning meta-framework.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree that this paper is a useful and valuable contributions to ML engineering.
 - insightful analysis .. highly user friendly operator design
 - "useful and I can see it having large adoption in the community of scientific computing" ... "
 - "Personally I tend to buy these advantages of einops" ... "However, there is a lack of solid empirical study to validate the effectiveness and efficiency of the design"
 - "a useful and appealing new coding tool."

The negative reviewers appear fixated on the (true) observation that the paper does not look like a conventional ICLR paper, thati it "reads like a technical blog", and "lacks rigour".
I belive it is fair and measured to state that these reviews may be considered to exhibit aspects of gatekeeping: requiring more "mathiness" that does not help the paper, or more "rigour" through user studies that are in fact less valuable than the reviewers' own observations "I could see myself...", "I tend to buy...".

This is a paper about design, not about models or algorithms (although the algorithmic work is good).  It is about the design of tools that we all use, and about the decisions and thought processes that led to that design.  A reviewer decries "many non-rigorous claims".  These are claims about the usability of existing systems, and mostly appear in the discussion and footnotes, as the authors note in rebuttal.  Of course, one could have run user studies to back up each claim, but I am just as convinced by the examples shown in the paper.  It matters not to me what some users corralled into a user study thought.  It matters what I and my colleagues will think, and I am now sure to recommend einops to colleagues.  I would not have met it had the paper not been submitted to ICLR, and hence I am certain it should be accepted, so more can see that we care not just about mathiness, but actually enabling progress in our field.

The job of a conference like ICLR is to expose researchers and practitioners in machine learning to ideas and techniques that may advance their research and practice.  Programming, and the translation of mathematical ideas to efficient computer code, are fundamental to all of machine learning, and hence programming models are very much suitable for presentation to an ICLR audience.  I am certain that this paper, and the technology it describes, are more important to ICLR readers than knowing that if module A is co-trained with module B, then combined with compression C, the SOTA on some arbitrary benchmark is increased by 0.31 +/- 1.04. 

Reviewer gRMH says "there is no code", but the code has been in the open for three years; it is an accident of our misapplication of the principles of blind review that the reviewer felt they could not search for the code, and that the authors felt they could not bring to bear the evidence that three years of real-world usage have brought.  

Reviewers say the work is just an extension of einsum, while noting that the extension is useful and nontrivial.  Yes, it is an extension, and the paper's examples show how it yields more compact code that is also more readable and maintainable.

I could add more examples, but in short, I tend to side with the authors' response at almost every point.  At the same time, the final version of the paper has been strengthened by this dialectic, and I expect further strenghtening through exposure to the ICLR community.

To the authors: Listing 1 is useful, but should be in an appendix.  Instead, add examples of ellipses on P5, and show more inline examples in general.  The paper would be strengthened by another pass over the English -- after the decision is made I would be happy to volunteer to help.
Decision: Accept (Oral)