Tsanley: Understanding Tensor Programs
Tsanley is a shape analyzer for tensor programs, using popular tensor libraries:
numpy. Plugs into your existing code seamlessly, with minimal changes.
Builds upon the library tsalib for specifying, annotating and transforming tensor shapes using named dimensions.
tsanley discovers shape errors at runtime by checking the runtime tensor shapes against the user-specified shape annotations. Tensor shape annotations are specified in the
tsalib shape shorthand notation, e.g.,
More details on the shorthand format here.
Suppose we have the following functions
test_foo in our existing code. To setup
tsanley analyzer for shape checking in
foo, we add a function
setup_named_dims before calling
test_foo, label tensor variables by their expected shorthand shapes (e.g.,
b,d) and then execute the code normally.
def foo(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] def test_foo(): import torch x = torch.Tensor(10, 100, 1024) foo(x) def setup_named_dims(): from tsalib import dim_vars #declare the named dimension variables using the tsalib api #e.g., 'b' stands for 'Batch' dimension with size 10 dim_vars('Batch(b):10 Length(t):100 Hidden(d):1024') # initialize tsanley's dynamic shape analyzer from tsanley.dynamic import init_analyzer init_analyzer(trace_func_names=['foo'], show_updates=True) #check_tsa=True, debug=False if __name__ == '__main__': setup_named_dims() test_foo()
On executing the above program,
tsanley tracks shapes of tensor variables (
z) in function
foo and reports following shape check results.
> Analyzing function foo Update at line 36: actual shape of x = b,t,d >> shape check succeeded at line 36 Update at line 37: actual shape of y = t,d >> FAILED shape check at line 37 expected: (b:10, d:1024), actual: (100, 1024) Update at line 38: actual shape of z = b,d >> shape check succeeded at line 38 saving shapes to /tmp/shape_log.json ..
What does setup_named_dims do?
- Declare the named dimension variables (using
dim_vars) -- using them we can specify the expected shape of tensor variables in the code. For example, here we declare 3 dimension variables,
Hidden, and refer to them via shorthand names
- We use shorthand names to label tensor variables and check their shapes in one or more functions, e.g.,
- Initialize the
tsanleyanalyzer by calling
trace_func_namestakes a list of function names as Unix shell-style wildcards (using the
fnmatchlibrary). We can specify names with wildcards, e.g.,
Resnet.*to track all functions in the
See examples in models directory.
pip install tsanley
tsanley can also annotate tensor variables in existing executable code 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 file
- First, generate shape logs by adding
model.py. The logs are stored in
- Use the logs to annotate
Let's revisit the earlier example, without our manual annotations. Suppose it resides in
def foo(x): y = x.mean(dim=0) z = x.mean(dim=1) def test_foo(): import torch x = torch.Tensor(10, 100, 1024) foo(x)
setup_named_dims to the code, and execute it.
def setup_named_dims(): from tsalib import dim_vars #declare the named dimension variables using the tsalib api #e.g., 'b' stands for 'Batch' dimension with size 10 dim_vars('Batch(b):10 Length(t):100 Hidden(d):1024') # initialize tsanley's dynamic shape analyzer from tsanley.dynamic import init_analyzer init_analyzer(trace_func_names=['foo'], show_updates=True, check_tsa=False) # debug=False if __name__ == '__main__': setup_named_dims() test_foo()
This generates the shape logs in
check_tsa=False ensures no shape checks are performed by
foo with the command:
tsa annotate -f model.py
The output is a file
foo updated as follows:
def foo(x): y: 't,d' = x.mean(dim=0) z: 'b,d' = x.mean(dim=1)
tsanley makes smart guesses to map runtime shape values (
100) to the shorthand names (
t). If we do not declare the dimension names using
setup_named_dims, we get the following annotation:
def foo(x): y: '100,1024' = x.mean(dim=0) z: '10,1024' = x.mean(dim=1)
tsanley performs 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.
numpy programs should also work (
tsalib supported backends), but remain to be tested.
Release history Release notifications | RSS feed
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.