Package: tsanley
tsanleyTsanley is a shape analyzer for tensor programs, using popular tensor libraries:tensorflow,pytorch,numpy. Plugs into your existing code seamlessly, with minimal changes.Builds upon the librarytsalibfor specifying, annotating and transforming tensor shapes usingnamed dimensions.Quick Starttsanleydiscovers shape errors at runtime by checking the runtime tensor shapes against the user-specified shape annotations. Tensor shape annotations are specified in thetsalibshape shorthand notation, e.g.,x: 'btd'.More details on the shorthand formathere.ExampleSuppose we have the following functionsfooandtest_fooin our existing code. To setuptsanleyanalyzer for shape checking infoo, we add a functionsetup_named_dimsbeforecallingtest_foo, label tensor variables by their expected shorthand shapes (e.g.,b,d) and then execute the code normally.deffoo(x):x:'b,t,d'#shape check: ok!               [line 36]y:'b,d'=x.mean(dim=0)# error!         [line 37]z:'b,d'=x.mean(dim=1)#shape check: ok! [line 38]deftest_foo():importtorchx=torch.Tensor(10,100,1024)foo(x)defsetup_named_dims():fromtsalibimportdim_vars#declare the named dimension variables using the tsalib api#e.g., 'b' stands for 'Batch' dimension with size 10dim_vars('Batch(b):10 Length(t):100 Hidden(d):1024')# initialize tsanley's dynamic shape analyzerfromtsanley.dynamicimportinit_analyzerinit_analyzer(trace_func_names=['foo'],show_updates=True)#check_tsa=True, debug=Falseif__name__=='__main__':setup_named_dims()test_foo()On executing the above program,tsanleytracks shapes of tensor variables (x,y,z) in functionfooand reports following shape check results.Output>AnalyzingfunctionfooUpdateatline36:actualshapeofx=b,t,d>>shapechecksucceededatline36Updateatline37:actualshapeofy=t,d>>FAILEDshapecheckatline37expected:(b:10,d:1024),actual:(100,1024)Updateatline38:actualshapeofz=b,d>>shapechecksucceededatline38savingshapesto/tmp/shape_log.json..What does setup_named_dims do?Declare the named dimension variables (usingdim_vars) -- using them we can specify the expected shape of tensor variables in the code. For example, here we declare 3 dimension variables,Batch,LengthandHidden, and refer to them via shorthand namesb,t,d.We use shorthand names to label tensor variables and check their shapes in one or more functions, e.g.,foohere.Initialize thetsanleyanalyzer by callinginit_analyzer: parametertrace_func_namestakes a list of function names as Unix shell-style wildcards (using thefnmatchlibrary). We can specify names with wildcards, e.g.,Resnet.*to track all functions in theResnetclass.See examples inmodelsdirectory.Installationpip install tsanleyAnnotationtsanleycan also annotate tensor variables in existingexecutablecode with shape labels. This is useful when trying to understand external open-source code or labeling one's own code.Suppose, we have some un-annotated code residing in filemodel.py.First, generateshape logsby addingsetup_named_dimsto themodel.py.Executemodel.py. The logs are stored in/tmp/shape_log.json.Use the logs to annotatetest.py.ExampleLet's revisit the earlier example, without our manual annotations. Suppose it resides inmodel.py.deffoo(x):y=x.mean(dim=0)z=x.mean(dim=1)deftest_foo():importtorchx=torch.Tensor(10,100,1024)foo(x)We addsetup_named_dimsto the code, and execute it.defsetup_named_dims():fromtsalibimportdim_vars#declare the named dimension variables using the tsalib api#e.g., 'b' stands for 'Batch' dimension with size 10dim_vars('Batch(b):10 Length(t):100 Hidden(d):1024')# initialize tsanley's dynamic shape analyzerfromtsanley.dynamicimportinit_analyzerinit_analyzer(trace_func_names=['foo'],show_updates=True,check_tsa=False)# debug=Falseif__name__=='__main__':setup_named_dims()test_foo()This generates the shape logs in/tmp/shape_log.json. Flagcheck_tsa=Falseensures no shape checks are performed bytsanley.Now, annotatefoowith the command:tsa annotate -f model.pyThe output is a filetsa_model.pywithfooupdated as follows:deffoo(x):y:'t,d'=x.mean(dim=0)z:'b,d'=x.mean(dim=1)tsanleymakes smart guesses to map runtime shape values (100) to the shorthand names (t). If we do not declare the dimension names usingdim_varsinsetup_named_dims, we get the following annotation:deffoo(x):y:'100,1024'=x.mean(dim=0)z:'10,1024'=x.mean(dim=1)Status: Experimentaltsanleyperforms a best-effort shape tracking during program execution. Here are a few tricky scenarios:calling same function multiple times -- shape values from only the last call are cached.recursive calls -- not handled.Tested withpytorchexamples.tensorflowandnumpyprograms should also work (tsalibsupported backends), but remain to be tested.