Package: torchtyping
torchtypingType annotations for a tensor's shape, dtype, names, ...Turn this:defbatch_outer_product(x:torch.Tensor,y:torch.Tensor)->torch.Tensor:# x has shape (batch, x_channels)# y has shape (batch, y_channels)# return has shape (batch, x_channels, y_channels)returnx.unsqueeze(-1)*y.unsqueeze(-2)into this:defbatch_outer_product(x:TensorType["batch","x_channels"],y:TensorType["batch","y_channels"])->TensorType["batch","x_channels","y_channels"]:returnx.unsqueeze(-1)*y.unsqueeze(-2)with programmatic checking that the shape (dtype, ...) specification is met.Bye-bye bugs! Say hello to enforced, clear documentation of your code.If (like me) you find yourself littering your code with comments like# x has shape (batch, hidden_state)or statements likeassert x.shape == y.shape, just to keep track of what shape everything is,then this is for you.InstallationpipinstalltorchtypingRequires Python 3.7+ and PyTorch 1.7.0+.Usagetorchtypingallows for type annotating:shape: size, number of dimensions;dtype(float, integer, etc.);layout(dense, sparse);namesof dimensions as pernamed tensors;arbitrary number of batch dimensionswith...;...plus anything else you like, astorchtypingis highly extensible.Iftypeguardis (optionally) installed thenat runtime the types can be checkedto ensure that the tensors really are of the advertised shape, dtype, etc.# EXAMPLEfromtorchimportrandfromtorchtypingimportTensorType,patch_typeguardfromtypeguardimporttypecheckedpatch_typeguard()# use before @typechecked@typecheckeddeffunc(x:TensorType["batch"],y:TensorType["batch"])->TensorType["batch"]:returnx+yfunc(rand(3),rand(3))# worksfunc(rand(3),rand(1))# TypeError: Dimension 'batch' of inconsistent size. Got both 1 and 3.typeguardalso has an import hook that can be used to automatically test an entire module, without needing to manually add@typeguard.typecheckeddecorators.If you're not usingtypeguardthentorchtyping.patch_typeguard()can be omitted altogether, andtorchtypingjust used for documentation purposes. If you're not already usingtypeguardfor your regular Python programming, then strongly consider using it. It's a great way to squash bugs. Bothtypeguardandtorchtypingalso integrate withpytest, so if you're concerned about any performance penalty then they can be enabled during tests only.APItorchtyping.TensorType[shape,dtype,layout,details]The core of the library.Each ofshape,dtype,layout,detailsare optional.Theshapeargument can be any of:Anint: the dimension must be of exactly this size. If it is-1then any size is allowed.Astr: the size of the dimension passed at runtime will be bound to this name, and all tensors checked that the sizes are consistent.A...: An arbitrary number of dimensions of any sizes.Astr: intpair (technically it's a slice), combining bothstrandintbehaviour. (Just astron its own is equivalent tostr: -1.)Astr: strpair, in which case the size of the dimension passed at runtime will be bound tobothnames, and all dimensions with either name must have the same size. (Some people like to use this as a way to associate multiple names with a dimension, for extra documentation purposes.)Astr: ...pair, in which case the multiple dimensions corresponding to...will be bound to the name specified bystr, and again checked for consistency between arguments.None, which when used in conjunction withis_namedbelow, indicates a dimension that mustnothave a name in the sense ofnamed tensors.ANone: intpair, combining bothNoneandintbehaviour. (Just aNoneon its own is equivalent toNone: -1.)ANone: strpair, combining bothNoneandstrbehaviour. (That is, it must not have a named dimension, but must be of a size consistent with other uses of the string.)Atyping.Any: Any size is allowed for this dimension (equivalent to-1).Any tuple of the above. For example.TensorType["batch": ..., "length": 10, "channels", -1]. If you just want to specify the number of dimensions then use for exampleTensorType[-1, -1, -1]for a three-dimensional tensor.Thedtypeargument can be any of:torch.float32,torch.float64etc.int,bool,float, which are converted to their corresponding PyTorch types.floatis specifically interpreted astorch.get_default_dtype(), which is usuallyfloat32.Thelayoutargument can be eithertorch.stridedortorch.sparse_coo, for dense and sparse tensors respectively.Thedetailsargument offers a way to pass an arbitrary number of additional flags that customise and extendtorchtyping. Two flags are built-in by default.torchtyping.is_namedcauses thenames of tensor dimensionsto be checked, andtorchtyping.is_floatcan be used to check that arbitrary floating point types are passed in. (Rather than just a specific one as with e.g.TensorType[torch.float32].) For discussion on how to customisetorchtypingwith your owndetails, see thefurther documentation.Check multiple things at once by just putting them all together inside a single[]. For exampleTensorType["batch": ..., "length", "channels", float, is_named].torchtyping.patch_typeguard()torchtypingintegrates withtypeguardto perform runtime type checking.torchtyping.patch_typeguard()should be called at the global level, and will patchtypeguardto checkTensorTypes.This function is safe to run multiple times. (It does nothing after the first run).If using@typeguard.typechecked, thentorchtyping.patch_typeguard()should be called any time before using@typeguard.typechecked. For example you could call it at the start of each file usingtorchtyping.If usingtypeguard.importhook.install_import_hook, thentorchtyping.patch_typeguard()should be called any time before defining the functions you want checked. For example you could calltorchtyping.patch_typeguard()just once, at the same time as thetypeguardimport hook. (The order of the hook and the patch doesn't matter.)If you're not usingtypeguardthentorchtyping.patch_typeguard()can be omitted altogether, andtorchtypingjust used for documentation purposes.pytest--torchtyping-patch-typeguardtorchtypingoffers apytestplugin to automatically runtorchtyping.patch_typeguard()before your tests.pytestwill automatically discover the plugin, you just need to pass the--torchtyping-patch-typeguardflag to enable it. Packages can then be passed totypeguardas normal, either by using@typeguard.typechecked,typeguard's import hook, or thepytestflag--typeguard-packages="your_package_here".Further documentationSee thefurther documentationfor:FAQ;Includingflake8andmypycompatibility;How to write custom extensions totorchtyping;Resources and links to other libraries and materials on this topic;More examples.