Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation

1 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 . Available research-and-production frameworks for DL , such as pytorch ( Paszke et al. , 2019 ) , tensorflow ( Abadi et al. , 2015 ) , mxnet ( Chen et al. , 2015 ) , jax ( Bradbury et al. , 2018 ) , and others , vary in numerous aspects , but their core functionality is built around efficient computations on n-dimensional arrays ( tensors for brevity1 ) . 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 . Due to growing usage of DL in production and rising complexity of models , it 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 or 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 a 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 to resolve mentioned problems in related works . Our approach , einops – a verbose notation for tensor manipulation – is introduced next , followed by the code examples in case studies . We implement the notation in a Python package , and explain main design choices in implementation details section . We conclude with a discussion of einops role and common criticisms . ∗Currently at Herophilus , Inc. 1Our sincere apologies to readers with backgrounds in mathematics and physics for possible confusion . 2 MAINSTREAM APPROACH FOR TENSOR OPERATIONS . Nowadays tensor programming is dominating in DL and playing a crucial role in 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 on the following assumptions : • Tensor is an n-dimensional array with shape ( e.g . 4× 6× 3 ) and data type ( e.g . float32 ) . • When elements are matched across tensors , axes 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 to be either BHWC or BCHW ordered ) or should be provided with indices of axes that should be treated separately:2 y = x.max ( axis=1 ) # some operations are provided with indices y = x.swap_axes ( 0 , 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 : a 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 a chain of manipulations a 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 # im_stack has shape [ b , 3 , h , w ] y1 = im_stack.transpose ( 2 , 0 , 3 , 1 ) .reshape ( h , b * w , 3 ) y2 = im_stack.transpose ( 2 , 3 , 0 , 1 ) .reshape ( h , b * w , 3 ) One of y1 , y2 contains a correct image , while the other is irrelevant – it requires some time to 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 to “ 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 can not meaningfully visualize intermediate layers , so there is no way to narrow down searches for a problem source . Thus , a researcher/engineer has to check all the code after a failure . Traditional operations restrict code flexibility : any change in the shape agreements between parts of the code is hard to align and all related code fragments ( frequently located in different places ) should be updated simultaneously . In the next code fragment we add omitted batch dimension to the code from the vision permutator ( Hou et al. , 2021 ) , and then update code to support depth : # pytorch-like code without batch dimension , as in the paper x_h = x.reshape ( H , W , N , S ) .permute ( 2 , 1 , 0 , 3 ) .reshape ( N , W , H * S ) x_h = proj_h ( x_h ) .reshape ( N , W , H , S ) .permute ( 2 , 1 , 0 , 3 ) .reshape ( H , W , C ) # with batch dimension x_h = x.reshape ( B , H , W , N , S ) .permute ( 0 , 3 , 2 , 1 , 4 ) .reshape ( B , N , W , H * S ) x_h = proj_h ( x_h ) .reshape ( B , N , W , H , S ) .permute ( 0 , 3 , 2 , 1 , 4 ) .reshape ( B , H , W , C ) # with batch and depth dimension x_h = x.reshape ( B , H , W , D , N , S ) .permute ( 0 , 4 , 2 , 3 , 1 , 5 ) .reshape ( B , N , W , D , H * S ) x_h = proj_h ( x_h ) .reshape ( B , N , W , D , H , S ) .permute ( 0 , 4 , 2 , 3 , 1 , 5 ) .reshape ( B , H , W , D , C ) 2Pseudocode in the paper corresponds to numpy unless otherwise stated . There is no way to write crossframework code . This problem we partially address with proposed einops notation . Modifications are very error-prone : all indices should be recomputed and order of reshape operands should be verified . Uncommonly , this fragment is quite self-documenting since final reshape in each line hints the intended order of axes after the transposition . In DL code it is common to use x.transpose ( 0 , 3 , 1 , 2 ) expecting other users to recognize a familiar pattern or leaving a comment.3 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 0 3 1 ) is inverse to ( 1 3 0 2 ) . Less critical , but still notable source of mistakes is 0-based enumeration of axes ( Julia , MATLAB and R use 1-based ) , while we propose the framework not relying on axis numeration at 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 between human communication and code . Finally , common tensor operations can easily break the tensor structure . For example , reshape , a common operation in the DL code , easily breaks the tensor structure because a whole tensor is treated as a sequence and no connection is assumed between axes in input and output , see Figure 1 . Aforementioned implementation mistakes can result in rejecting valuable research or drawing incorrect 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 an error . Even the presence of an 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 . A number of recently proposed architectures ( Wu et al. , 2021 ; Tolstikhin et 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 . 3 RELATED WORKS . A commonly taken approach to increase reliability is assigning names/labels to tensor dimensions ( below we refer to them as labeled tensors ) . Most known implementations are xarray ( Hoyer & Hamman , 2017 ) , labeled_tensor in tensorflow ( Christiansen & Hoyer , 2016 ) , namedtensor ( 2019 ) , and named tensors in Pytorch ( 2019 ) . Despite being around for many years , this idea got little adoption , and not used in the DL 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 . : # x1 has axes ( x , y , height ) and x2 has axes ( time , x , y ) x1 + x2 # result has axes ( x , y , height , time ) , maybe in a different order x1.mean ( 'height ' ) # reduce over axis named 'height' 3Out of 50 last usages of torch.permute in Python on GitHub on 22 Nov 2021 , only in 4 cases code had comments about result shape , in other cases no information was provided in any other form , including reshape . Axes matching rules vary across implementations . However , we can describe some common issues why this approach did not gain wide support in the DL community : • Operations focus on modified axes , and neither describe input nor output ; a user has to remember axes labels for each of 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 , e.g . for namedtensor package : # x1 has axes ( height , width ) # x2 has axes ( h , w ) # x3 has axes ( heigt ) x1 + x2 + x3 # has axes ( height width h w heigt ) in some order • Adoption requires a strong buy-in decision , as all code should be axis label-centric . In contrary , transition to 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 , a type/class/interface/protocol is introduced to describe a 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 a language model can expect to manipulate sequences of letters and thus expects axis “ letter ” . However , surrounding code may try to use it for prediction of words , pixels or phonemes . Thus , relabeling of axes may be required to “ connect ” subsystems . In 2011 , M. Wiebe introduced an operation einsum in numpy package . With some simplifications ( absence of covariant and contravariant indices , contracted dimension may be not repeated ) einsum mimics Einstein summation rule commonly used in physics . numpy.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 to other packages : tensorflow , pytorch , mxnet , chainer Tokui et al . ( 2019 ) , etc . numpy.einsum ( 'ij , jk- > ik ' , A , B ) # matrix multiplication numpy.einsum ( 'ijk- > ij ' , C ) # sum over last axis numpy.einsum ( 'ij , ji- > ' , A , B ) # trace of matrix product einops revisits and expands this approach with an emphasis on complex tensor layouts and rearrangements , additional checks and broader functionality . In our work we try to align interface with einsum to allow smooth simultaneous usage . However , interface adjustments ( such as 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 einops Einstein-Inspired Notation for OPerationS , einops , is our proposal to address problems mentioned in Section 2 . The core of our approach is a new notation for transformation patterns and , in its basic view , this notation enumerates the tensor elements in one-to-one-correspondence 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 new axes ) • all axis identifiers on either side of expression should be unique Examples of einops transformation patterns are ' b c h w - > b h w c ' # transpose ' b c h w - > b c ' # reduce on h , w ' b c - > b c h w ' # repeat on h , w ' ( h1 h2 ) ( w1 w2 ) c - > ( h1 w1 ) h2 w2 c ' # 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 one tensor pattern to be different ( thus traces , permitted by einsum , are not provided by einops ) . This also requires that ellipsis can be used only once within a 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 in a number : # x is of shape ( 10 , 10 , 10 , 10 ) , then x [ 6 , 2 , 4 , 9 ] == y [ 6249 ] y = rearrange ( x , ' a b c d - > ( a b c 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 i3 - > ( i1 i2 i3 ) # Let original array have a shape of [ 2 , 3 , 2 ] , result has a length 2x3x2=12 i1 0 0 0 0 0 0 1 1 1 1 1 1 i2 0 0 1 1 2 2 0 0 1 1 2 2 i3 0 1 0 1 0 1 0 1 0 1 0 1 final position 0 1 2 3 4 5 6 7 8 9 10 11 Reverse pattern ( i1 i2 i3 ) - > i1 i2 i3 uses the same bijection to decompose an axis into three . Since axis can be decomposed in multiple ways ( e.g . 12 could be represented as 2 × 3 × 2 or 1× 12× 1 or 3× 1× 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 a view of the original data . For example , rearrangement ' ( a b c ) ( d e f ) ( g h ) - > a b ( c d ) e ( f g 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 a special meaning – they do not correspond to an axis variable , and thus their behavior is separate from anonymous axes . ' h w - > h w 1 ' # add unitary axis ' h w - > h w 3 ' # repeat values along new anonymous axis ' 1 h w 3 - > h w ' # remove unitary axis and reduce on anonymous axis of length 3 ' ... h w - > ... w h ' # transpose two last dimensions ' b ... c - > ( ... ) b c ' # compose all but first and last dimensions # and move resulting new axis to front ' b c ... - > b c ' # reduce on all but first two dimensions All anonymous axes are treated as different even if they share length . Similar to einsum , ellipsis works as a wildcard for zero or more axes , their number and lengths are derived from input shape ( s ) . Ellipses are not named ( grammar allows adding names later ) , thus all ellipses within a 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 , q , v of shape [ batch , seq , channel ] , one can turn it into multi-head attention for 3-dimensional data by composing height , width and depth to a single dimension , and grouping head and batch dimension to ensure independent processing of attention heads : b h w d head c - > ( b head ) ( h w d ) c. Likewise , other neural blocks can be “ upgraded ” by rearranging inputs and outputs . 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 ( im , ' ( b1 b2 ) h w c - > ( b1 h ) ( b2 w ) c ' , b1=4 , b2=4 ) # max-pooling with kernel of size 2x2 reduce ( im , ' b c ( h h2 ) ( w w2 ) - > b c h w ' , 'max ' , h2=2 , w2=2 ) # 2x upsampling of individual image by repeating pixels repeat ( im , ' h w c - > ( 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 a wrong pattern is passed.4 In einsum , when an axis is present in all tensors , operation performs independently for all values of this axis , which is principle in einops . A tensor pattern only identifies correspondence between axes and element ’ 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 is checked if provided . • A tensor structure can not be broken by design , because the notation connects input axes ( 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 ) operations5 in Listing 1 1 np.transpose ( x , [ 0 , 3 , 1 , 2 ] ) rearrange ( x , ' b h w c - > b c h w ' ) 2 np.reshape ( x , rearrange ( x , ' h w c - > ( h w ) c ' ) 3 [ x.shape [ 0 ] * x.shape [ 1 ] , x.shape [ 2 ] ] ) 4 np.squeeze ( x , 0 ) rearrange ( x , ' ( ) h w c - > h w c ' ) 5 np.expand_dims ( x , -1 ) rearrange ( x , ' h w c - > h w c ( ) ' ) 6 np.stack ( [ r , g , b ] , axis=2 ) rearrange ( [ r , g , b ] , ' c h w - > h w c ' ) 7 np.concatenate ( [ r , g , b ] , axis=0 ) rearrange ( [ r , g , b ] , ' c h w - > ( c h ) w ' ) 4Post-factum we can confirm this choice : search over Github shows that rearrange , the most restrictive on possible patterns , also accounts for the majority of usages . 5Some examples use list inputs , details see in Appendix B . 8 np.flatten ( x ) rearrange ( x , ' b t c - > ( b t c ) ' ) 9 np.swap_axes ( x , 0 , 1 ) rearrange ( x , ' b t c - > t b c ' ) 10 left , right = np.split ( image , 2 , axis=1 ) rearrange ( x , ' h ( lr w ) c - > lr h w c ' , lr=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 h w c - > b c ' , 'max ' ) 13 np.mean ( x ) reduce ( x , ' b h w c - > ' , 'mean ' ) 14 np.mean ( x , axis= ( 1 , 2 ) , keepdims=True ) reduce ( x , ' b h w c - > b ( ) ( ) c ' , 'mean ' ) 15 np.reshape ( x , [ -1 , 2 ] ) .max ( axis=1 ) reduce ( x , ' ( h 2 ) - > h ' , 'max ' ) 16 np.repeat ( x , 2 , axis=1 ) repeat ( x , ' h w - > h ( w 2 ) ' ) 17 np.tile ( x , 2 , axis=1 ) repeat ( x , ' h w - > h ( 2 w ) ' ) 18 np.tile ( x [ : , : , np.newaxis ] , 3 , axis=2 ) repeat ( x , ' h w - > h w 3 ' ) Listing 1 : Correspondence between numpy and einops operations .

Summary: Operators are key elements in deep learning (TensorFlow, PyTorch, etc) or scientific computing (NumPy) frameworks. This paper introduces a generic, convenient and elegant way of massively describing large portion of frequently used operators under a unified framework, EinOps, inspired by einsum, or the Einstein summation convention. Under this convention, key components like multi-head attention in modern neural networks can be expressed with fewer lines of code. The major contribution of this work includes: - Presents a set of novel and highly expressive notations to represent hundreds of operators in modern deep learning frameworks; - Develops a library that interprets the proposed notations and offloads them to existing implementations in NumPy or TensorFlow, PyTorch; - Demonstrates significantly improved readability and debuggability with the proposed notations