Package: einops
https://user-images.githubusercontent.com/6318811/177030658-66f0eb5d-e136-44d8-99c9-86ae298ead5b.mp4einopsFlexible and powerful tensor operations for readable and reliable code.Supports numpy, pytorch, tensorflow, jax, andothers.Recent updates:0.7.0: no-hassletorch.compile, support ofarray api standardand more10'000🎉: github reports that more than 10k project use einopssee how to use einops withtorch.compileeinops 0.6.1: paddle backend addedeinops 0.6 introducespacking and unpackingeinops 0.5: einsum is now a part of einopsEinops paperis accepted for oral presentation at ICLR 2022 (yes, it worth reading).
Talk recordings areavailablePrevious updates- flax and oneflow backend added
- torch.jit.script is supported for pytorch layers
- powerful EinMix added to einops. [Einmix tutorial notebook](https://github.com/arogozhnikov/einops/blob/master/docs/3-einmix-layer.ipynb)TweetsIn case you need convincing arguments for setting aside time to learn about einsum and einops...Tim Rocktäschel, FAIRWriting better code with PyTorch and einops 👌Andrej Karpathy, AI at TeslaSlowly but surely, einops is seeping in to every nook and cranny of my code. If you find yourself shuffling around bazillion dimensional tensors, this might change your lifeNasim Rahaman, MILA (Montreal)More testimonialsContentsInstallationDocumentationTutorialAPI micro-referenceWhy using einopsSupported frameworksCitingRepositoryanddiscussionsInstallationPlain and simple:pipinstalleinopsTutorialsTutorials are the most convenient way to seeeinopsin actionpart 1:einops fundamentalspart 2:einops for deep learningpart 3:packing and unpackingpart 4:improve pytorch code with einopsKapil Sachdeva recorded a smallintro to einops.APIeinopshas a minimalistic yet powerful API.Three core operations provided (einops tutorialshows those cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate, view and numerous reductions)fromeinopsimportrearrange,reduce,repeat# rearrange elements according to the patternoutput_tensor=rearrange(input_tensor,'t b c -> b c t')# combine rearrangement and reductionoutput_tensor=reduce(input_tensor,'b c (h h2) (w w2) -> b h w c','mean',h2=2,w2=2)# copy along a new axisoutput_tensor=repeat(input_tensor,'h w -> h w c',c=3)Later additions to the family arepackandunpackfunctions (better than stack/split/concatenate):fromeinopsimportpack,unpack# pack and unpack allow reversibly 'packing' multiple tensors into one.# Packed tensors may be of different dimensionality:packed,ps=pack([class_token_bc,image_tokens_bhwc,text_tokens_btc],'b * c')class_emb_bc,image_emb_bhwc,text_emb_btc=unpack(transformer(packed),ps,'b * c')Finally, einops provides einsum with a support of multi-lettered names:fromeinopsimporteinsum,pack,unpack# einsum is like ... einsum, generic and flexible dot-product# but 1) axes can be multi-lettered  2) pattern goes last 3) works with multiple frameworksC=einsum(A,B,'b t1 head c, b t2 head c -> b head t1 t2')EinMixEinMixis a generic linear layer, perfect for MLP Mixers and similar architectures.LayersEinops provides layers (einopskeeps a separate version for each framework) that reflect corresponding functionsfromeinops.layers.torchimportRearrange,Reducefromeinops.layers.tensorflowimportRearrange,Reducefromeinops.layers.flaximportRearrange,Reducefromeinops.layers.paddleimportRearrange,Reducefromeinops.layers.kerasimportRearrange,Reducefromeinops.layers.chainerimportRearrange,ReduceExample of using layers within a pytorch modelExample given for pytorch, but code in other frameworks is almost identicalfromtorch.nnimportSequential,Conv2d,MaxPool2d,Linear,ReLUfromeinops.layers.torchimportRearrangemodel=Sequential(...,Conv2d(6,16,kernel_size=5),MaxPool2d(kernel_size=2),# flattening without need to write forwardRearrange('b c h w -> b (c h w)'),Linear(16*5*5,120),ReLU(),Linear(120,10),)No more flatten needed!Additionally, torch users will benefit from layers as those are script-able and compile-able.Namingeinopsstands for Einstein-Inspired Notation for operations
(though "Einstein operations" is more attractive and easier to remember).Notation was loosely inspired by Einstein summation (in particular bynumpy.einsumoperation).Why useeinopsnotation?!Semantic information (being verbose in expectations)y=x.view(x.shape[0],-1)y=rearrange(x,'b c h w -> b (c h w)')While these two lines are doing the same job insomecontext,
the second one provides information about the input and output.
In other words,einopsfocuses on interface:what is the input and output, nothowthe output is computed.The next operation looks similar:y=rearrange(x,'time c h w -> time (c h w)')but it gives the reader a hint:
this is not an independent batch of images we are processing,
but rather a sequence (video).Semantic information makes the code easier to read and maintain.Convenient checksReconsider the same example:y=x.view(x.shape[0],-1)# x: (batch, 256, 19, 19)y=rearrange(x,'b c h w -> b (c h w)')The second line checks that the input has four dimensions,
but you can also specify particular dimensions.
That's opposed to just writing comments about shapes since comments don't prevent mistakes, not tested, and without code review tend to be outdatedy=x.view(x.shape[0],-1)# x: (batch, 256, 19, 19)y=rearrange(x,'b c h w -> b (c h w)',c=256,h=19,w=19)Result is strictly determinedBelow we have at least two ways to define the depth-to-space operation# depth-to-spacerearrange(x,'b c (h h2) (w w2) -> b (c h2 w2) h w',h2=2,w2=2)rearrange(x,'b c (h h2) (w w2) -> b (h2 w2 c) h w',h2=2,w2=2)There are at least four more ways to do it. Which one is used by the framework?These details are ignored, sinceusuallyit makes no difference,
but it can make a big difference (e.g. if you use grouped convolutions in the next stage),
and you'd like to specify this in your code.Uniformityreduce(x,'b c (x dx) -> b c x','max',dx=2)reduce(x,'b c (x dx) (y dy) -> b c x y','max',dx=2,dy=3)reduce(x,'b c (x dx) (y dy) (z dz) -> b c x y z','max',dx=2,dy=3,dz=4)These examples demonstrated that we don't use separate operations for 1d/2d/3d pooling,
those are all defined in a uniform way.Space-to-depth and depth-to space are defined in many frameworks but how about width-to-height? Here you go:rearrange(x,'b c h (w w2) -> b c (h w2) w',w2=2)Framework independent behaviorEven simple functions are defined differently by different frameworksy=x.flatten()# or flatten(x)Supposex's shape was(3, 4, 5), thenyhas shape ...numpy, pytorch, cupy, chainer:(60,)keras, tensorflow.layers, gluon:(3, 20)einopsworks the same way in all frameworks.Independence of framework terminologyExample:tilevsrepeatcauses lots of confusion. To copy image along width:np.tile(image,(1,2))# in numpyimage.repeat(1,2)# pytorch's repeat ~ numpy's tileWith einops you don't need to decipher which axis was repeated:repeat(image,'h w -> h (tile w)',tile=2)# in numpyrepeat(image,'h w -> h (tile w)',tile=2)# in pytorchrepeat(image,'h w -> h (tile w)',tile=2)# in tfrepeat(image,'h w -> h (tile w)',tile=2)# in jaxrepeat(image,'h w -> h (tile w)',tile=2)# in cupy...(etc.)Testimonialsprovide users' perspective on the same question.Supported frameworksEinops works with ...numpypytorchtensorflowjaxcupychainertf.kerasoneflow(experimental)flax(experimental)paddle(experimental)Additionally, starting from einops 0.7.0 einops can be used with any framework that supportsPython array API standardCiting einopsPlease use the following bibtex record@inproceedings{
    rogozhnikov2022einops,
    title={Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation},
    author={Alex Rogozhnikov},
    booktitle={International Conference on Learning Representations},
    year={2022},
    url={https://openreview.net/forum?id=oapKSVM2bcj}
}Supported python versionseinopsworks with python 3.8 or later.