Package: tensor-annotations
TensorAnnotationsTensorAnnotations is an experimental library enabling annotation of data-type
and semantic shape information using type annotations - for example:defcalculate_loss(frames:Array4[uint8,Time,Batch,Height,Width]):...This annotation states that the data-type offramesisuint8, and that the
dimensions are time-like, batch-like, etc. (while saying nothing about the
actualvalues- e.g. the actual batch size).Why? Two reasons:Shape annotations can be checkedstatically. This can catch a range of
bugs caused by e.g. wrong selection or reduction of axes before you run your
code - even when the errors would not necessarily throw a runtime exception!Interface documentation (also enabling shape autocompletion in IDEs).To do this, the library provides three things:A set of custom tensor types for TensorFlow and JAX, supporting the above
kinds of annotationsA collection of common semantic labels (e.g.Time,Batch, etc.)Type stubs for common library functions that preserve semantic shape
information (e.g.reduce_sum(Tensor[Time, Batch], axis=0) -> Tensor[Batch])TensorAnnotations is being developed for JAX and TensorFlow.ExampleHere is some code that takes advantage of static shape checking:importtensorflowastffromtensor_annotationsimportaxesimporttensor_annotations.tensorflowasttfuint8=ttf.uint8Batch,Time=axes.Batch,axes.Timedefsample_batch()->ttf.Tensor2[uint8,Time,Batch]:returntf.zeros((3,5))deftrain_batch(batch:ttf.Tensor2[uint8,Batch,Time]):m:ttf.Tensor1[uint8,Batch]=tf.reduce_max(batch,axis=1)# Do something usefuldefmain():batch1=sample_batch()batch2=tf.transpose(batch1)train_batch(batch2)This code contains shape annotations in the signatures ofsample_batchandtrain_batch, and in the line callingreduce_max. It is otherwise the
same code you would have written in an unchecked program.You can check these annotations for inconsistencies by running a static type
checker on your code (see 'General usage' below). For example, runningtrain_batchdirectly onbatch1will result in the following error from
pytype:File "example.py", line 10: Function train_batch was called with the wrong arguments [wrong-arg-types]
         Expected: (batch: Tensor2[uint8, Batch, Time])
  Actually passed: (batch: Tensor2[uint8, Time, Batch])Similarly, changing the the call toreduce_maxfromaxis=1toaxis=0results in:File "example.py", line 15: Type annotation for m does not match type of assignment [annotation-type-mismatch]
  Annotation: Tensor1[uint8, Batch]
  Assignment: Tensor1[uint8, Time](These messages were shortened for readability. The actual errors will be more
verbose because fully qualified type names will be displayed. We are looking
into improving this.)Seeexamples/tf_time_batch.pyfor a complete example.RequirementsTensorAnnotatations requires Python 3.8 or above, due to the use oftyping.Literal.InstallationTo install custom tensor types:pipinstalltensor_annotationsThen, depending on whether you use JAX or TensorFlow:pipinstalltensor_annotations_jax_stubs# and/orpipinstalltensor_annotations_tensorflow_stubsIf you use pytype, you'll also need to take a few extra steps to let
it take advantage of JAX/TensorFlow stubs (since it doesn't yet support
PEP 561 stub packages). First, make a copy of typeshed in e.g. your home
directory:gitclonehttps://github.com/python/typeshed"$HOME/typeshed"Next, symlink the stubs into your copy of typeshed:site_packages=$(python3-msite--user-site)# Custom tensor classesmkdir-p"$HOME"/typeshed/stubs/{tensor_annotations/tensor_annotations,tensorflow,jax}ln-s"$site_packages/tensor_annotations/__init__.py""$HOME/typeshed/stubs/tensor_annotations/tensor_annotations/__init__.pyi"ln-s"$site_packages/tensor_annotations/jax.pyi""$HOME/typeshed/stubs/tensor_annotations/tensor_annotations/jax.pyi"ln-s"$site_packages/tensor_annotations/tensorflow.pyi""$HOME/typeshed/stubs/tensor_annotations/tensor_annotations/tensorflow.pyi"ln-s"$site_packages/tensor_annotations/axes.py""$HOME/typeshed/stubs/tensor_annotations/tensor_annotations/axes.pyi"# TensorFlowln-s"$site_packages/tensorflow-stubs""$HOME/typeshed/stubs/tensorflow/tensorflow"# JAXln-s"$site_packages/jax-stubs""$HOME/typeshed/stubs/jax/jax"General usageFirst, importtensor_annotationsand start annotating function signatures
and variable assignments. This can be done gradually.Next, run a static type checker on your code. If you use Mypy, it should just
work. If you use pytype, you need to invoke it in a special way in order to let
it know about the custom typeshed installation:TYPESHED_HOME="$HOME/typeshed" pytype your_code.pyWe recommend you deliberately introduce a shape error and then confirm that
your type checker gives you an error to be sure you're set up correctly.Annotated tensor classesTensorAnnotations provides tensor classes for JAX and TensorFlow:# JAXimporttensor_annotations.jaxastjaxtjax.arrayN# Where N is the rank of the tensor# TensorFlowimporttensor_annotations.tensorflowasttfttf.TensorN# Where N is the rank of the tensorThese classes can be parameterized by semantic axis labels (below) using
generics, similar toList[int]. (Different classes are needed for each rank
because Python currently does not support variadic generics, but we're working
on it.)Data typesTensorAnnotations also provides its own data-type types:# JAXfromtensor_annotations.jaximportuint8,float32# Etc# TensorFlowfromtensor_annotations.tensorflowimportuint8,float32# EtcThis is because, for various reasons, the native data-type types liketf.uint8andjnp.uint8are unsuitable for use in type annotations. Seetensorflow.pyandjax.pyfor more information.Axis labelsAxis labels are used to indicate the semantic meaning of each dimension in a
tensor - whether the dimension is batch-like, features-like, etc. Note that no
connection is made between the symbol, e.g.Batch, and the actualvalueof
that dimension (e.g. the batch size) - the symbol really does only describe the
semantic meaning of the dimension.Seeaxes.pyfor the list of axis labels we provide out of the box. To define a
custom axis label, simply subclasstensor_annotations.axes.Axis. You can also
usetyping.NewTypeto do this using a single line:CustomAxis=typing.NewType('CustomAxis',axes.Axis)In the future we intend to support axis types that are tied to the actual size
of that axis. Currently, however, we don't have a good way of doing this. If you
nonetheless want to annotate certain dimensions with a literal size, e.g. for
documentation of interfaces which are hardcoded for specific sizes, we recommend
you just use a custom axis for this purpose. (Just to be clear, though: these
sizes willnotbe checked - neither statically, nor at runtime!)L64=typing.NewType('L64',axes.Axis)StubsBy default, TensorFlow and JAX are not aware of our annotations. For example, if
you have a tensorx: Array2[uint8, Time, Batch]and you calljnp.sum(x, axis=0), you won't get aArray1[uint8, Batch], you'll just get anAny. We therefore provide a set of custom type annotations for TensorFlow and
JAX packaged in 'stub' (.pyi) files.Our stubs currently cover the following parts of the API. All operations are
supported for rank 1, 2, 3 and 4 tensors, unless otherwise noted. Unary
operators are also supported for rank 0 (scalar) tensors.TensorFlowSeeCoverage.Tensor unary operators: For tensorx:abs(x),-x,+xTensor binary operators: For tensorsaandb:a + b,a / b,a // b,a ** b,a < b,a > b,a <= b,a >= b,a * b. Yet to be typed:a ? float,a ? intforTensor0, broadcasting where one axis is 1JAXSeeCoverage.Tensor unary operators: For tensorx,abs(x),-x,+xTensor binary operators: For tensorsaandb,a + b,a / b,a // b,a ** b,a < b,a > b,a <= b,a >= b,a * b. Yet to be typed:a ? float,a ? intforTensor0, broadcasting where one axis is 1CastingSome of your code might be already typed with existing library tensor types:defsample_batch()->jnp.array:...If this is the case, and you don't want to change these types globally in your
code, you can cast to TensorAnnotations classes withtyping.cast:fromtypingimportcastx=cast(tjax.Array2[uint8,Batch,Time],x)Note that this is only a hint to the type checker - at runtime, it's a no-op. An
alternative syntax emphasising this fact is:x:tjax.Array2[uint8,Batch,Time]=x# type: ignoreGotchasUse tuples for shape/axis specificationsFor type inference with TensorFlow and JAX API functions we often have to match
additional arguments. I.e., the rank of atf.zeros(...)tensor depends on the
length of the shape argument. This only works with tuples, and not with lists:a=tf.zeros((10,10))# Correctly infers type Tensor2[Any, Any]b:ttf.Tensor2[uint8,Time,Batch]=get_batch()c=tf.transpose(b,perm=(0,1))# Tracks and infers the axes-types of bwhilea=tf.zeros([10,10])# Returns Anyb:ttf.Tensor2[uint8,Time,Batch]=get_batch()c=tf.transpose(b,perm=[0,1]))# Does not track permutations and returns AnyRuntime vs static checksNote that we do not verify that the rank of a tensor at runtime matches the one
specified in the annotations. If you were in an evil mood, you could create an
untyped (Any) tensor, and statically type it as something completely wrong. This
is in line with the rest of the python type-checking approach, which does notenforceconsistency with the annotated types at runtime.Value consistency. Not only do we not verify the rank, we don't verify
anything about the actual shape value either. The following willnotraise an
error:x:tjax.Array1[uint8,Batch]=jnp.zeros((3,))y:tjax.Array1[uint8,Batch]=jnp.zeros((5,))Note thatthis is by design! Shape symbols such asBatcharenotplaceholders for actual values like 3 or 5. Symbols only refer to thesemantic
meaningof a dimension. In the above example, say,xmight be a train batch,
andymight be a test batch, and therefore they have different sizes, even
though both of their dimensions are batch-like. This means that even
element-wise operations likez = x + ywould in this case not raise a
type-check error.FAQsWhy doesn't e.g.tjax.ArrayNsubclassjnp.DeviceArray?We'dlikethis to be the case, but haven't figured out how to yet because
of circular dependencies:ArrayNis defined intensor_annotations.jax, which would need to importjax.numpyin order to subclassjnp.DeviceArray.However, ourjax.numpystubs make use ofArrayN, sojax.numpyitself
needs to importtensor_annotations.jax.We ultimate solution to this will hopefully be to upstream ourArrayNclasses
such that they can be defined injax.numpytoo. Until then, we'll just be
trying to make e.g.tjax.ArrayNlook as close tojnp.DeviceArrayas possible
through dummy methods and dummy attributes so that autocomplete still works.
If there are particular methods/attributes you'd like added, please do let us
know.Why are so many methods annotated asAnyin the JAX stubs?We don't yet have a good way of automatically generating stubs in general.
For the methods where wedogenerate stubs automatically (all the ones
not annotated asAny), we've checked their signature manually and written
stub generators for each method individually.Ideally we'd start from stubs generated by e.g. pytype and then customise
them to include shape information, but we haven't got around to setting
this up yet.Why not usePEP 646?Compatibility. There are two factors: a) concise syntax for PEP 646 is
only available in Python 3.11 onwards, which not everyone can migrate to yet;
and b) PEP 646 is (as of January 2022) only supported by Pyre and Pyright -
not by Mypy or pytype, which are both popular.Type checker support is the biggest thing - so once thereisbetter support
for PEP 646 in Mypy and pytype, we may revisit this question.See alsoThis library is one approach of many to checking tensor shapes. We don't expect
it to be the final solution; we create it to explore one point in the space of
possibilities.Other tools for checking tensor shapes include:Pythia, a static analyzer
designed specifically for detecting TensorFlow shape errorstsanley, which uses string annotations
combined with runtime verificationPyContracts, a general-purpose
library for specifying constraints on function arguments that has special
support for NumPyShape Guard, another runtime
verification tool using concise helper methodsswift-tfp, a static analyzer
for tensor shapes in SwiftTo learn more about tensor shape checking in general, see:Stephan Hoyer'sIdeas for array shape typing in PythondocumentTheTyping for multi-dimensional arraysGitHub issue inpython/typingOurShape annotation feature scopingand ourShape annotation syntax proposaldocuments (a synthesis of the most promising ideas
from the full doc)The Pythontyping-sigmailing list (in particular,this thread)Notes and recordingsfrom the Tensor Typing Open Design MeetingsRepository structureThetensor_annotationspackage contains four types of things:Custom tensor classes. We provide our own versions of e.g. TensorFlow'sTensorclass and JAX'sArrayclass in order to support shape
parameterisation. These are stored intensorflow.pyandjax.py.
(Note that these are only used in the context of type annotations - they are
never instantiated - hence no implementation being present.)Type stubs for custom tensor classes. We also need to provide type
annotations specifying what the shape of, say,x: Tensor[A, B] + y: Tensor[B]is. These aretensorflow.pyiandjax.pyi.These are generated fromtemplates/tensors.pyiusingtools/render_tensor_template.py.Type stubs for library functions. Finally, we need to specify what the
shape of, say,tf.reduce_sum(x: Tensor[A, B], axis=0)is. This information
is stored in type stubs inlibrary_stubs. (Thethird_party/pydirectory structure is necessary to indicate to pytype exactly which packages
these stubs are for.) Ideally, these will eventually live in the libraries
themselves.JAX stubs are auto-generated fromtemplates/jax.pyiusingtools/render_jax_library_template.pyi. Note that we currently specify the
signature of the library members we don't generate automatically asAny.
Ideally, we'd like to automatically generate complete type stubs and then
tweak them to include shape information, but we haven't gotten around to
this yet.For TensorFlow stubs, we start from stubs generated by a Google-internal
TensorFlow stub generator
and then hand-edit those stubs to include shape stubs. The edits we've made
are demarcated byBEGIN/END tensor_annotations annotations for ...blocks.
Again, we'll make this more automated in the future.Common axis types. Finally, we also provide a canonical set of common axis
labels such as 'time' and 'batch'. These are stored inaxes.py.