Tsanley: Understanding Tensor Programs
Project description
tsanley
A shape analyzer for tensor programs, using popular tensor libraries: tensorflow
, pytorch
, numpy
, ...
Builds upon a library tsalib for specifying, annotating and transforming tensor shapes using named dimensions.
Quick Start
tsanley
discovers shape errors at runtime by checking the runtime tensor shapes with the user-specified shape annotations. Tensor shape annotations are specified in the tsalib
shape shorthand notation, e.g., x: 'btd'
.
More details on the shorthand format here.
Example
Suppose we have the following functions foo
and test_foo
in our existing code. To setup tsanley
, we add the following code before test_foo
is called (in the test
function):
- Declare the named dimension variables (using
dim_vars
) -- these help specify the expected shape of tensor variables used in the code. For example, here we declare 3 dimension variables, referred to via shorthand namesb
,t
,d
. We use these shorthand names to label variables and check their shapes in one or more functions, e.g.,foo
here. - Initialize the
tsanley
analyzer by callinginit_analyzer
: parametertrace_func_names
takes a list of function names as Unix shell-style wildcards (using thefnmatch
library).
def foo(x):
x: 'b,t,d' #shape check: ok! [line 36]
y: 'b,d' = x.mean(dim=0) #error: dim should be 1 [line 37]
# ^
# | tsanley detects shape violation
z: 'b,d' = x.mean(dim=1) #shape check: ok! [line 38]
def test_foo():
import torch
from tsalib import get_dim_vars
# get the declared dimension sizes: 10, 100, 1024
B, L, D = get_dim_vars('b t d')
x = torch.Tensor(B, L, D)
foo(x)
def test():
#declare the named dimension variables using the tsalib api
from tsalib import dim_vars
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
test_foo()
if __name__ == '__main__': test()
On executing the above program, tsanley
tracks shapes of tensor variables (x
, y
, z
) in function foo
and reports shape check successes and failures.
Output
> 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: (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 ..
See examples in models directory.
Installation
pip install tsanley
Status: Experimental
tsanley
performs a best-effort shape tracking when the program runs. 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 with pytorch
examples. tensorflow
and numpy
programs should also work (tsalib
supported backends).
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.