TSAlib: Support for Named Tensor Shapes
Project description
Tensor Shape Annotations Library (tsalib)
Conventional tensor manipulation libraries — numpy
, pytorch
, keras
, tensorflow
, lack support for naming the dimensions of tensor variables. tsalib
enables using named dimensions with existing libraries, using Python's support for type annotations and a new shorthand notation for naming tensor shapes (TSN).
Why named dimensions ? See References.
Updates:
Using tsalib
:

track shapes: label your tensor variables with their named shapes (
x: 'btd'
orx: (B,T,D)
) 
better debugging: write named shape assertions (
assert x.shape == (B,T,D)
). 
write seamless named shape transformations:
warp(x, '(btd)* > btdl > bdtl > b,d//2,t*2,l', 'jpv')
instead of a sequence of calls over a laundry list of APIs (
reshape
,permute
,stack
,concat
) 
work with arbitrary backends without changes:
numpy
,pytorch
,keras
,tensorflow
,mxnet
, etc.
Exposing the invisible named dimensions enhances code clarity, accelerates debugging and leads to improved productivity across the board. Even complex deep learning architectures need only a small number of named dimensions.
The complete API in a notebook here, an introductory article here.
Contents
 Quick Start
 Installation
 Design Principles, Model Examples (includes BERT!)
 API Overview
 Best Practices  How to use
tsalib
 Change Log
Quick Start
from tsalib import dim_vars as dvs, size_assert import tensorflow as tf import torch #declare dimension variables (from config arguments) B, C, H, W = dvs('Batch(b):32 Channels(c):3 Height(h):256 Width(w):256') ... # create tensors using dimension variables (interpret dim vars as integers) x: 'bchw' = torch.randn(B, C, H, W) x: 'bchw' = tf.get_variable("x", shape=(B, C, H, W), initializer=tf.random_normal_initializer()) # perform tensor transformations, keep track of named shapes x: 'b,c,h//2,w//2' = maxpool(x) # check assertions: compare dynamic shapes with declared shapes # assertions are 'symbolic': don't change even if declared shapes change assert x.size() == (B, C, H // 2, W // 2) #or, check selected dimensions size_assert (x.size(), (B,C,H//2,W//2), dims=[1,2,3])
Write intuitive and crisp shape transformations:
# A powerful onestop `warp` operator to compose multiple transforms inline # here: a sequence of a permute ('p') and view ('v') transformations y = warp(x1, 'bhwc > bchw > b*c,h,w', 'pv') assert y.size() == (B*C,H,W) #or, the same transformation sequence with anonymous dims y = warp (x1, ['_hwc > _chw', 'bc,, > b*c,,'], 'pv') # Combinations of `alignto`, `dot` and broadcast # Enables writing really compact code for similar patterns ht: 'bd'; Wh: 'dd'; Y: 'bld'; WY: 'dd' a: 'bd' = dot('_d.d_', ht, Wh) b: 'b,1,d' = alignto((a,'bd'), 'bld') Mt: 'bld' = torch.tanh(dot('__d.d_', Y, WY) + b)
Old vs New code
def merge_heads_old(x): x = x.permute(0, 2, 1, 3).contiguous() new_x_shape = x.size()[:2] + (x.size(2) * x.size(1),) res = x.view(*new_x_shape)
def merge_heads_tsalib(x: 'bhtd'): res: 'b,t,h*d' = warp(x, 'bhtd > bthd > b,t,h*d', 'pcv')
Named shapes may be represented as tuples or shorthand strings. Details on shorthand notation here.
 a tuple
(B,H,D)
[long form] or a string'b,h,d'
(or simply'bhd'
) [shorthand]  a string with anonymous dimensions (
',h,'
or_h_
is a 3d tensor).
Installation
pip install [upgrade] tsalib
Documentation, Design Principles, Model Examples
This notebook serves as a working documentation for the tsalib
library and illustrates the complete tsalib
API. The shorthand notation is documented here.
The models directory contains tsalib annotations of a few wellknown, complex neural architectures:
 BERT.
 OpenAI Transformer,
 Resnet,
 Contrast models with and without tsalib (pytorch, tensorflow).
With named shape annotations, we can gain deeper and immediate insight into how the module works by scanning through the forward
(or equivalent) function.
tsalib
is designed to stay light and easy to incorporate into existing workflow with minimal code changes. Choose to usetsalib
for tensor labels and shape asserts only, or, integrate deeply by usingwarp
everywhere in your code. The API includes both libraryindependent and dependent parts, giving developers flexibility in how they choose to incorporate
tsalib
in their workflow.  We've carefully avoided deeper integration into popular tensor libraries to keep
tsalib
lightweight and avoid backendinflicted bugs.
API
from tsalib import dim_vars as dvs, get_dim_vars import numpy as np
Declarations
Declare, Update Dimension Variables
#or declare dim vars with default integer values (optional) B, C, D, H, W = dvs('Batch:48 Channels:3 EmbedDim:300 Height Width') #or provide *shorthand* names and default values for dim vars [best practice] B, C, D, H, W = dvs('Batch(b):48 Channels(c):3 EmbedDim(d):300 Height(h) Width(w)') # switch from using config arguments to named dimensions B, C, D = dvs('Batch(b):{0} Channels(c):{1} EmbedDim(d):{2}'.format(config.batch_size, config.num_channels, config.embed_dim))
Update dimension variable length dynamically.
H.update_len(1024) print(f'H = {H}') # h:1024 update_dim_vars_len({'h': 512, 'w': 128}) # ... H, W = get_dim_vars('h w') print(f'H, W = {H}, {W}') # h:512, w:128
Annotate Variables with Shapes
Annotate with shorthand named shapes.
a: 'b,d' = np.array([[1., 2., 3.], [10., 9., 8.]]) #(Batch, EmbedDim): (2, 3) b: '2bd' = np.stack([a, a]) #(2, Batch, EmbedDim): (2, 2, 3)
Annotations are optional and do not affect program performance. Arithmetic over dimension variables is supported. This enables easy tracking of shape changes across neural network layers.
B, C, H, W = get_dim_vars('b c h w') #lookup predeclared dim vars v: 'bchw' = torch.randn(B, C, h, w) x : 'b,c*2,h//2,w//2' = torch.nn.conv2D(C, C*2, ...)(v)
Shape and Tensor Transformations
Onestop shape transforms: warp
operator
The warp
operator enables squeezing in a sequence of shape transformations in a single line using shorthand notation (TSN). warp
takes in an input tensor, a sequence of shape transformations, and the corresponding transform types (view transform > 'v', permute transform > 'p').
x: 'btd' = torch.randn(B, T, D) y = warp(x, 'btd > b,t,4,d//4 > b,4,t,d//4 ', 'vp') #(v)iew, then (p)ermute, transform assert(y.shape == (B,4,T,D//4))
Because it returns transformed tensors, the warp
operator is backend librarydependent. Currently supported backends are numpy
, tensorflow
and pytorch
. New backends can be added easily (see backend.py). See docs for transform types here.
Or, perform individual transforms using named shapes.
#use dimension variables directly x = torch.ones(B, T, D) x = x.view(B, T, 4, D//4) from tsalib import view_transform as vt, permute_transform as pt y = x.reshape(vt('btd > b,t,4,d//4', x.shape)) #(20, 10, 300) > (20, 10, 4, 75) assert y.shape == (B, T, 4, D//4) y = x.transpose(pt('b,,d, > d,,b,'))
See notebook for more examples.
More useful operators: join
, alignto
, reduce_dims
...
more ..
Unified stack/concat
using join
. Join together sequence of tensors into a single tensor in different ways using the same join
operator. join
is backenddependent.
# xi : (B, T, D) # "concatenate" along the 'T' dimension: "(b,t,d)* > (b,3*t,d)" x = tsalib.join([x1, x2, x3], ',*,') assert x.shape == (B, 3*T, D) # "stack": join by adding a new dimension to the front: "(b,t,d)* > (^,b,t,d)" x = join([x1, x2, x3], '^') assert x.shape == (3, B, T, D)
Align one tensor to the rank of another tensor using alignto
.
x1 = np.random.randn(D,D) x2 = np.random.randn(B,D,T,D) x1_aligned = alignto((x1, 'dd'), 'bdtd') assert x1_aligned.shape == (1,D,1,D) x1_aligned = alignto((x1, 'dd'), 'bdtd', tile=True) assert x1_aligned.shape == (B,D,T,D)
Use dimension names instead of cryptic indices in reduction (mean
, max
, ...) operations.
from tsalib import reduce_dims as rd b: (2, B, D) c: (D,) = np.mean(b, axis=rd('2bd > d')) #axis = (0,1)
Simplified dot
operator
Easy matmult
specification when
 exactly a single dimension is common between the operands and
 the order of dimensions preserved in the output.
x = torch.randn(B, C, T) y = torch.randn(C, D) z = dot('_c_.c_', x, y) assert z.size() == (B, T, D)
Dependencies
sympy
. A library for building symbolic expressions in Python is the only dependency.
Tested with Python 3.6, 3.7.
For writing type annotations inline, Python >= 3.5 is required which allows optional type annotations for variables. These annotations do not affect the program performance in any way.
Best Practices
tsalib
is designed for progressive adoption with your current deep learning models and pipelines. You can start off only with declaring and labeling variables with named shapes and writing shape assertions. This already brings tremendous improvement in productivity and code readability. Once comfortable, use other transformations:warp
,join
, etc. Convert all relevant config parameters into dimension variables. Use only latter in your code.
 Define most (if not all) dimension variables upfront  this requires some discipline. Use
get_dim_vars
to lookup predefined dimension variables by their shorthand names in any function context. Update dimension variables dynamically in code withupdate_dim_vars_len
.  Avoid using
reshape
: useview
andtranspose
together. An inadvertentreshape
may not preserve your dimensions (axes). Usingview
to change shape protects against this: it throws an error if the dimensions being manipulated are not contiguous.  Shape Annotations vs Assertions. Shape labels (
x: (B,T,D)
orx: 'btd'
) ease shape recall during coding. Shape assertions (assert x.shape === (B,T,D)
) enable catching inadvertent shape bugs at runtime. Pick either or both to work with.
References
 Blog article introducing
tsalib
 A proposal for designing a tensor library with named dimensions from groundup.
tsalib
takes care of some use cases, without requiring any change in the tensor libraries.  Pytorch Issue on Names Axes here.
 Using einsum for tensor operations improves productivity and code readability. blog
 The Tile DSL uses indices ranging over dimension variables to write compact, libraryindependent tensor operations.
 The datashape library introduces a generic type system and grammar for structure data.
tsalib
focuses on shapes of homogeneous tensor data types only, with arithmetic support.  The xarray library.
 The einops library.
 The NamedTensor library.
 The TensorNetwork library. Generalizes the idea of named axes and composition/decomposition/reordering of axes very nicely.
Writing deep learning programs which manipulate multidim tensors (numpy
, pytorch
, tensorflow
, ...) requires you to carefully keep track of shapes of tensors. In absence of a principled way to name tensor dimensions and track shapes, most developers resort to writing adhoc shape comments embedded in code (see code from googleresearch/bert) or spaghetti code with numeric indices: x.view(* (x.size()[:2] + (x.size(2) * x.size(1),))
. This makes both reading — figuring out RNN
output shapes, examining/modifying deep pretrained architectures (resnet
, densenet
, elmo
) — and writing — designing new kinds of attention
mechanisms (multihead attention
)— deep learning programs harder.
Update: Pytorch has recently introduced support for named shapes in tensors  naming is optional and lazy, like in tsalib
. This is great news! We hope that tensorflow and numpy (in particular!) will also incorporate named shapes, to make it fundamentally easier for the developer community to write tensor programs .
tsalib
’s USP continues to be the shorthand notation for writing shape transformations and more importantly, annotating tensor variables. See discussion here. Using shorthands exploits the fact that most dimension names remain the same during the program execution and reduces overhead and program clutter due to shape names. We hope that pytorch and other libraries will also incorporate shorthand shape naming in the near future.
Developers benefit from shape annotations/assertions in many ways:
Benefits: * Quickly verify the variable shapes when writing new transformations or modifying existing modules. * Assertions and annotations remain the same even if the actual dimension sizes change. * Faster *debugging*: if you annotateasyougo, the tensor variable shapes are explicit in code, readily available for a quick inspection. No more adhoc shape `print`ing when investigating obscure shape errors. * Do shape transformations using *shorthand* notation and avoid unwanted shape surgeries. * Use named shapes to improve code clarity everywhere, even in your machine learning data pipelines. * They serve as useful documentation to help others understand or extend your module.Author
 Nishant Sinha, OffNote Labs (nishant@offnote.co, @medium, @twitter)
Change Log
The library is stable. Contributions/feedback welcome!
 [May 2020] Added
update_dim_vars_len
andDimExpr.update_len
.  [5 Feb 2019] Added
dot
operator.  [4 Feb 2019] Added fully annotated and adapted BERT model. More illustrative pytorch and tensorflow snippets.
 [31 Jan 2019] Added
alignto
operator.  [18 Dec 2018] Added the
join
operator.warp
takes a list of (shorthand) transformations.  [28 Nov 2018] Added
get_dim_vars
to lookup dim vars declared earlier. Shorthand notation docs.  [21 Nov 2018] Added documentation notebook.
 [18 Nov 2018] Support for
warp
,reduce_dims
. Backend modules fornumpy
,tensorflow
andtorch
added.  [9 Nov 2018] Support for shorthand notation in view/permute/expand transforms.
 [9 Nov 2018] Support for using named shapes in assertions and tensor constructors (cast to integers).
 [25 Oct 2018] Initial Release
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.
Filename, size  File type  Python version  Upload date  Hashes 

Filename, size tsalib0.2.2py3noneany.whl (24.9 kB)  File type Wheel  Python version py3  Upload date  Hashes View 
Filename, size tsalib0.2.2.tar.gz (19.6 kB)  File type Source  Python version None  Upload date  Hashes View 