Package: tsalib
Tensor Shape Annotations Library (tsalib)Conventional tensor manipulation libraries —numpy,pytorch,keras,tensorflow, lack support fornamingthe dimensions of tensor variables.tsalibenables usingnamed dimensionswith existing libraries, using Python's support fortype annotationsand a newshorthand notationfor naming tensor shapes (TSN).Whynamed dimensions ?  SeeReferences.Updates:An in-depthtalkon tsalib at Anthill 2019.A dynamic shape checker based on tsalib:tsanleyUsingtsalib:track shapes: label your tensor variables with theirnamed shapes(x: 'btd'orx: (B,T,D))better debugging: writenamedshapeassertions(assert x.shape == (B,T,D)).write seamlessnamedshapetransformations:warp(x, '(btd)* -> btdl -> bdtl -> b,d//2,t*2,l', 'jpv')instead of a sequence of calls over a laundry list of APIs (reshape,permute,stack,concat)work with arbitrary backends without changes:numpy,pytorch,keras,tensorflow,mxnet, etc.Exposing theinvisiblenamed dimensions enhances code clarity, accelerates debugging and leads to improved productivity across the board. Even complex deep learning architectures need only a small number of named dimensions.The completeAPIin anotebookhere, an introductoryarticlehere.ContentsQuick StartInstallationDesign Principles, Model Examples(includesBERT!)API OverviewBest Practices -- How to usetsalibChange LogQuick Startfromtsalibimportdim_varsasdvs,size_assertimporttensorflowastfimporttorch#declare dimension variables (from config arguments)B,C,H,W=dvs('Batch(b):32 Channels(c):3 Height(h):256 Width(w):256')...# create tensors using dimension variables (interpret dim vars as integers)x:'bchw'=torch.randn(B,C,H,W)x:'bchw'=tf.get_variable("x",shape=(B,C,H,W),initializer=tf.random_normal_initializer())# perform tensor transformations, keep track of named shapesx:'b,c,h//2,w//2'=maxpool(x)# check assertions: compare dynamic shapes with declared shapes# assertions are 'symbolic': don't change even if declared shapes changeassertx.size()==(B,C,H//2,W//2)#or, check selected dimensionssize_assert(x.size(),(B,C,H//2,W//2),dims=[1,2,3])Write intuitive and crisp shape transformations:# A powerful one-stop `warp` operator to compose multiple transforms inline# here: a sequence of a permute ('p') and view ('v') transformationsy=warp(x1,'bhwc -> bchw -> b*c,h,w','pv')asserty.size()==(B*C,H,W)#or, the same transformation sequence with anonymous dimsy=warp(x1,['_hwc -> _chw','bc,, -> b*c,,'],'pv')# Combinations of `alignto`, `dot` and broadcast# Enables writing really compact code for similar patternsht:'bd';Wh:'dd';Y:'bld';WY:'dd'a:'bd'=dot('_d.d_',ht,Wh)b:'b,1,d'=alignto((a,'bd'),'bld')Mt:'bld'=torch.tanh(dot('__d.d_',Y,WY)+b)Old vs New codedefmerge_heads_old(x):x=x.permute(0,2,1,3).contiguous()new_x_shape=x.size()[:-2]+(x.size(-2)*x.size(-1),)res=x.view(*new_x_shape)defmerge_heads_tsalib(x:'bhtd'):res:'b,t,h*d'=warp(x,'bhtd -> bthd -> b,t,h*d','pcv')Named shapes may be represented as tuples orshorthandstrings. Details on shorthand notationhere.a tuple(B,H,D)[long form] or a string'b,h,d'(or simply'bhd') [shorthand]a string with anonymous dimensions (',h,'or_h_is a 3-d tensor).Installationpip install [--upgrade] tsalibDocumentation, Design Principles, Model ExamplesThisnotebookserves as a working documentation for thetsaliblibrary and illustrates the completetsalibAPI. Theshorthandnotation is documentedhere.Themodelsdirectory contains tsalib annotations of a few well-known, complex neural architectures:BERT.OpenAI Transformer,Resnet,Contrast models with and without tsalib (pytorch,tensorflow).With named shape annotations, we can gain deeper and immediate insight into how the module works by scanning through theforward(or equivalent) function.tsalibis designed to stay light and easy to incorporate into existing workflow with minimal code changes. Choose to usetsalibfor tensor labels and shape asserts only, or, integrate deeply by usingwarpeverywhere in your code.The API includes both library-independent and dependent parts, giving developers flexibility in how they choose to incorporatetsalibin their workflow.We've carefully avoided deeper integration into popular tensor libraries to keeptsaliblight-weight and avoid backend-inflicted bugs.APIfromtsalibimportdim_varsasdvs,get_dim_varsimportnumpyasnpDeclarationsDeclare, Update Dimension Variables#or declare dim vars with default integer values (optional)B,C,D,H,W=dvs('Batch:48 Channels:3 EmbedDim:300 Height Width')#or provide *shorthand* names and default values for dim vars [best practice]B,C,D,H,W=dvs('Batch(b):48 Channels(c):3 EmbedDim(d):300 Height(h) Width(w)')# switch from using config arguments to named dimensionsB,C,D=dvs('Batch(b):{0}Channels(c):{1}EmbedDim(d):{2}'.format(config.batch_size,config.num_channels,config.embed_dim))Updatedimension variable lengthdynamically.H.update_len(1024)print(f'H ={H}')# h:1024update_dim_vars_len({'h':512,'w':128})# ...H,W=get_dim_vars('h w')print(f'H, W ={H},{W}')# h:512, w:128Annotate Variables with ShapesAnnotate with shorthand named shapes.a:'b,d'=np.array([[1.,2.,3.],[10.,9.,8.]])#(Batch, EmbedDim): (2, 3)b:'2bd'=np.stack([a,a])#(2, Batch, EmbedDim): (2, 2, 3)Annotations are optional and do not affect program performance. Arithmetic over dimension variables is supported. This enables easy tracking of shape changes across neural network layers.B,C,H,W=get_dim_vars('b c h w')#lookup pre-declared dim varsv:'bchw'=torch.randn(B,C,h,w)x:'b,c*2,h//2,w//2'=torch.nn.conv2D(C,C*2,...)(v)Shape and Tensor TransformationsOne-stop shape transforms:warpoperatorThewarpoperator enables squeezing in asequenceof shape transformations in a single line using shorthand notation (TSN).warptakes in an input tensor, a sequence of shape transformations, and the corresponding transform types (view transform -> 'v', permute transform -> 'p').x:'btd'=torch.randn(B,T,D)y=warp(x,'btd -> b,t,4,d//4 ->  b,4,t,d//4 ','vp')#(v)iew, then (p)ermute, transformassert(y.shape==(B,4,T,D//4))Because it returns transformed tensors, thewarpoperator is backend library-dependent. Currently supported backends arenumpy,tensorflowandpytorch. New backends can be added easily (seebackend.py). See docs for transform typeshere.Or, perform individual transforms using named shapes.#use dimension variables directlyx=torch.ones(B,T,D)x=x.view(B,T,4,D//4)fromtsalibimportview_transformasvt,permute_transformaspty=x.reshape(vt('btd -> b,t,4,d//4',x.shape))#(20, 10, 300) -> (20, 10, 4, 75)asserty.shape==(B,T,4,D//4)y=x.transpose(pt('b,,d, -> d,,b,'))Seenotebookfor more examples.More useful operators:join,alignto,reduce_dims...more ..Unifiedstack/concatusingjoin. Join together sequence of tensors into a single tensor in different ways using the samejoinoperator.joinis backend-dependent.# xi : (B, T, D)# "concatenate" along the 'T' dimension: "(b,t,d)* -> (b,3*t,d)"x=tsalib.join([x1,x2,x3],',*,')assertx.shape==(B,3*T,D)# "stack": join by adding a new dimension to the front: "(b,t,d)* -> (^,b,t,d)"x=join([x1,x2,x3],'^')assertx.shape==(3,B,T,D)Align one tensor to the rank of another tensor usingalignto.x1=np.random.randn(D,D)x2=np.random.randn(B,D,T,D)x1_aligned=alignto((x1,'dd'),'bdtd')assertx1_aligned.shape==(1,D,1,D)x1_aligned=alignto((x1,'dd'),'bdtd',tile=True)assertx1_aligned.shape==(B,D,T,D)Use dimension names instead of cryptic indices inreduction(mean,max, ...) operations.fromtsalibimportreduce_dimsasrdb:(2,B,D)c:(D,)=np.mean(b,axis=rd('2bd -> d'))#axis = (0,1)SimplifieddotoperatorEasymatmultspecification whenexactly a single dimension is common between the operands andthe order of dimensions preserved in the output.x=torch.randn(B,C,T)y=torch.randn(C,D)z=dot('_c_.c_',x,y)assertz.size()==(B,T,D)Dependenciessympy. A library for building symbolic expressions in Python is the only dependency.Tested with Python 3.6, 3.7.For writing type annotations inline, Python >= 3.5 is required which allows optional type annotations for variables. These annotations do not affect the program performance in any way.Best Practicestsalibis designed forprogressive adoptionwith your current deep learning models and pipelines. You can start off only withdeclaringandlabelingvariables with named shapes and writing shapeassertions. This already brings tremendous improvement in productivity and code readability. Once comfortable, use other transformations:warp,join, etc.Convert all relevantconfig parametersinto dimension variables. Use only latter in your code.Definemost(if not all) dimension variables upfront -- this requires some discipline. Useget_dim_varsto lookup pre-defined dimension variables by their shorthand names in any function context. Update dimension variables dynamically in code withupdate_dim_vars_len.Avoid usingreshape: useviewandtransposetogether. An inadvertentreshapemay not preserve your dimensions (axes). Usingviewto change shape protects against this: it throws an error if the dimensions being manipulated are not contiguous.ShapeAnnotationsvsAssertions. Shape labels (x: (B,T,D)orx: 'btd') ease shape recall during coding. Shape assertions (assert x.shape === (B,T,D)) enable catching inadvertent shape bugs at runtime. Pick either or both to work with.ReferencesBlogarticleintroducingtsalibAproposalfor designing a tensor library with named dimensions from ground-up.tsalibtakes care of some use cases, without requiring any change in the tensor libraries.Pytorch Issue on Names Axeshere.Usingeinsumfor tensor operations improves productivity and code readability.blogTheTileDSL uses indices ranging over dimension variables to write compact, library-independent tensor operations.Thedatashapelibrary introduces a generic type system and grammar for structure data.tsalibfocuses on shapes of homogeneous tensor data types only, with arithmetic support.Thexarraylibrary.Theeinopslibrary.TheNamedTensorlibrary.TheTensorNetworklibrary. Generalizes the idea of named axes and composition/decomposition/reordering of axes very nicely.Writing deep learning programs which manipulate multi-dim tensors (numpy,pytorch,tensorflow, ...) requires you to carefully keep track of shapes of tensors. In absence of a principled way tonametensor dimensions and track shapes, most developers resort to writing adhoc shape comments embedded in code (see code fromgoogle-research/bert) or spaghetti code with numeric indices:x.view(* (x.size()[:-2] + (x.size(-2) * x.size(-1),)). This makes both reading — figuring outRNNoutput shapes, examining/modifying deep pre-trained architectures (resnet,densenet,elmo) — and writing —  designing new kinds ofattentionmechanisms (multi-head attention)— deep learning programs harder.Update: Pytorch has recentlyintroducedsupport for named shapes in tensors -- naming isoptionalandlazy, like intsalib. This is great news! We hope that tensorflow and numpy (in particular!) will also incorporate named shapes, to make it fundamentally easier for the developer community to write tensor programs .tsalib’s USP continues to be theshorthand notationfor writing shape transformations and more importantly, annotating tensor variables.  See discussionhere. Using shorthands exploits the fact that most dimension names remain the same during the program execution and reduces overhead and program clutter due to shape names. We hope that pytorch and other libraries will also incorporate shorthand shape naming in the near future.Developers benefit from shape annotations/assertions in many ways:Benefits:
    * Quickly verify the variable shapes when writing new transformations or modifying existing modules. 
    * Assertions and annotations remain the same even if the actual dimension sizes change.
    * Faster *debugging*: if you annotate-as-you-go, the tensor variable shapes are explicit in code, readily available for a quick inspection. No more adhoc shape `print`ing when investigating obscure shape errors.
    * Do shape transformations using *shorthand* notation and avoid unwanted shape surgeries.
    * Use named shapes to improve code clarity everywhere, even in your machine learning data pipelines.
    * They serve as useful documentation to help others understand or extend your module.AuthorNishant Sinha,OffNote Labs(nishant@offnote.co, @medium, @twitter)Change LogThe library is stable. Contributions/feedback welcome![May 2020] Addedupdate_dim_vars_lenandDimExpr.update_len.[5 Feb 2019] Addeddotoperator.[4 Feb 2019] Added fully annotated and adapted BERTmodel. More illustrative pytorch and tensorflow snippets.[31 Jan 2019] Addedaligntooperator.[18 Dec 2018] Added thejoinoperator.warptakes a list of (shorthand) transformations.[28- Nov 2018] Addedget_dim_varsto lookup dim vars declared earlier. Shorthand notation docs.[21 Nov 2018] Added documentationnotebook.[18 Nov 2018] Support forwarp,reduce_dims. Backend modules fornumpy,tensorflowandtorchadded.[9 Nov 2018] Support for shorthand notation in view/permute/expand transforms.[9 Nov 2018] Support for using named shapes in assertions and tensor constructors (cast to integers).[25 Oct 2018] Initial Release