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: 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.