When writing deep learning programs, keeping track of tensor shapes and dealing with subtle tensor shape errors (implicit broadcasts!!) gets quite frustrating.

We’ve been working on a tool tsanley (pronounced ‘stanley’) to enable finding subtle shape errors in your deep learning code quickly and cheaply. The key idea is to label tensor variables with their expected shapes (e.g., x : 'b,t,d' = ...) and let tsanley perform shape validity checks at runtime automatically. Works with small and big tensor programs.

repository: https://github.com/ofnote/tsanley

Quick example:

python def foo(x): x: 'b,t,d' #expected shape of x is (B, T, D). y: 'b,d' = x.mean(dim=0) * 2 # error! z: 'b,d' = x.mean(dim=1) # ok return y, z Function foo contains tensor variables labeled with their named shapes using a shorthand notation. It has a subtle shape error in the assignment to y: we expect the shape of y to be (B,D), however mean got rid of the first, and not the second, dimension. Your tensor library (pytorch / tensorflow / ..) won’t flag this as an error: instead, we will get a weird shape inconsistency error somewhere downstream.

tsanley finds such unexpected bugs quickly at runtime: “` 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 “`

Writing these named shape annotations manually can also get tedious. tsanley can auto-annotate the tensor variables in your (or someone else’s) code, if the code is executable. This is especially useful when trying to dig deep into or adapt an existing code / library for your project.

The tool builds upon the tsalib library, which introduced a shorthand notation for labeling tensor variables with their named shapes, irrespective of the backend tensor library used.

We would love feedback on tsanley and hope it is useful for your coding/debugging workflow.