Befriending Tensors with tsalib library

tldr: tsalib is a library to define dimension names and named shape expressions for tensors. Allows shape labels on variables, shape assertions and intuitive shape transformations using names. Works with arbitrary tensor libraries. Explicit shape annotations accelerate debugging of Deep Learning programs and improve developer productivity, code readability.

Source code available at the github repository.

Update (Nov 2019): Checkout our library tsanley to annotate and check named shapes on the fly. With tsanley, you can avoid writing explicit shape assertions and automatically annotate third party deep learning code that you want to reuse.

Writing deep learning programs which manipulate tensors (e.g., using numpy, pytorch, tensorflow, keras ..) requires you to carefully keep track of shapes of tensor variables. Carrying around the shapes in your head gets increasingly hard as you write more complex programs. For example, when creating a new RNN cell or designing a new kind of attention mechanism or trying to do a surgery of non-trivial pre-trained architectures (resnet101, densenet). Unfortunately, there is no principled way of shape tracking inside code — most developers resort to writing adhoc comments embedded in code to keep track of tensor shapes.

For instance, consider this snippet from the recent popular BERT model implementation from Google. Note how each statement is labeled by its input and output shapes. It is virtually impossible to code and debug complex tensor programs without such explicit shape tracking.

Code from BERT implementation from Google.

Similarly, consider this CNN model from the tensorflow tutorial. Notice how each statement is labeled with its input and output shapes to assist developers. Because shapes may be dynamic, these annotations contain both parameters (batch_size) and constants (28).

As network architectures get complicated, these annotations become essential for preserving developer sanity as well for others to understand and extend the code. Unfortunately, there is no native support for these annotations and engineers continue to write adhoc comments to fill the gap. In fact, practitioners have been asking for better support since a long time.

Tensor Shape Annotations

The TSA library (tsalib) comes to your rescue here.

• The library allows you to define custom dimension variables (B=Batch, C=Channels, …) and tensor shape expressions over them.
• tsalib does not touch the internals of the tensor library your code relies on — works well with any tensor library (numpy , tensorflow, pytorch, and so on).
• Label tensor variables with their shape expressions directly in the code, as first-class type annotations.
• You can TSAs to directly build (lucid) shape transformations ( reshape, permute, axis indices) and succinct matrix/tensor operations— more on this later.

Here is simple tensor program using TSAs (code available here). Note how variables a , b and c are labeled with their shapes. The shape annotations, e.g., (2,B,D) can contain both constants as well as named dimension variables ( B, D). Notice how the code becomes more readable instantly — you can simply glance down the shape annotations and discover the shape transformations performed by the numpyfunctions (even if you are unfamiliar with the API).

Dimension Variables, Shape Expressions

The core idea behind the library is that of dimension variables and expressions over them.

• tsalib allows you to declare custom dimension variables for your architecture, use them in shape annotations and perform arithmetic on them like ordinary Python variables.
• Declarations include optional size defaults and shorthand names. Replace your config parameters ( batch_size, embed_dim ) by dim variables ( B, D ). Use only dim vars in all parts of your code.

TSAs turn out to be useful in many ways.

• They help us to quickly cross-check and debug the variable shapes when writing new transformations or modifying existing modules.
• TSAs serve as useful documentation to guide others trying to understand or extend your module.
• Having explicit shapes around is beneficial for several reasons (see below).
• TSAs do not affect the program performance or correctness in any way. You can annotate more or less as per your needs.

Example: Resnet with TSAs

Here is a trickier example, a resnet module from the pytorch torchvision library. This module is the basis of multiple resnet implementations: resnet{18,34,50,101,152}. Now, it is pretty hard to gain any insight into the architecture by looking at the __init__ and _make_layer modules, even if you intuitively know the architecture well. Consequently, you can’t make sense of the forward function either.

Resnet Implementation (Pytorch)

Now, let’s look at the annotated forward function. Without needing to know the details of various component layers, the TSAs allow you to discover the shape of the module’s expected input (image) x and how it gets transformed by each layer. If you are reusing this module, e.g., in SSD for object detection or U-Net for segmentation, you can read off the shapes of all intermediate layers at a glance.

Note how the ability to write arithmetic expressions over named dimensions really helps. Config variables (num_classes) can be used in shapes directly. Also, it becomes clear that the architecture is parameterized over the block expansion variable Ex (1 in smaller resnets, 4 in bigger ones).

Attention-based architectures, e.g., Transformer, also suffer mightily from the opacity problem. See tsalib/examples/openai_transformer.py (annotated starting from the AllenNLP’s implementation) to see how TSAs significantly impact the readability of the code. Again, by skimming through the forward module, we get a quick insight into the architecture. Dissecting a multi-head attention architecture is now a piece of cake!

Shape Transformations, Shorthand notation

Once we have dimension variables in the code, we can exploit them to further improve code productivity and clarity. We can specify view and permute transforms directly using dimension variables. For example, this is how merge_heads is implemented in original openai_transformer.py from the AllenNLP library.

Now, here is how we can write it compactly using tsalib. All shapes are explicit and no more code surgery. Shape annotations can be made more succinct by using their shorthand string representation instead of tuples, e.g., 'bcd' instead of (B,C,D). tsalib allows using shorthands to specify transformations. Read more here.

Having TSA-like native support in all tensor libraries (numpy, pytorch, tensorflow) would be extremely beneficial to the practitioners. However, that may involve a serious re-design. In contrast, TSAs are library-independent, light-weight and plug-and-play. You can start using them today in your deep learning programs!

Do give tsalib a try and send me feedback — reach me at twitter or email me directly.

Library Internals

Type annotations (or type hints) is a well-supported, but optional Python 3 feature. Support for annotation as well as type checking is improving as Python 3 improves and is being adopted widely. For example, the AllenNLP guys use type and shape annotations throughout their code and recommend it as best practice.

tsalib uses a wonderful symbolic expression library called sympy internally to build and transform expressions. sympy is its only dependency.

About me: I’m an independent computer science researcher, engineer and speaker who loves to distill and transform complex technology into consumable products. Have worked across academia, industry and startups. I help companies understand and maneuver through the complex, evolving AI space and build deep learning based solutions that maximize ROI. If you enjoyed this article, please do clap and post your comments. You can follow me and read my other articles here, find me on linkedin or email me directly.