From e7ec5ba6ee4c0a0e655217d18c2b0a4289431ffd Mon Sep 17 00:00:00 2001 From: Sidney Zhang Date: Thu, 24 Jan 2019 10:18:34 -0500 Subject: [PATCH] add docs about annotating variables --- docs/stable/_sources/jit.rst.txt | 35 ++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/docs/stable/_sources/jit.rst.txt b/docs/stable/_sources/jit.rst.txt index 2baa500dc202..9b7e272b16b5 100644 --- a/docs/stable/_sources/jit.rst.txt +++ b/docs/stable/_sources/jit.rst.txt @@ -184,8 +184,13 @@ Example:: # and type int in the false branch By default, all parameters to a Torch Script function are assumed to be Tensor -because this is the most common type used in modules. To specify that an -argument to a Torch Script function is another type, it is possible to use +because this is the most common type used in modules. + +There are 2 scenarios in which you might want to annotate a type: + +1. Annotating Function Argument Types + +To specify that an argument to a Torch Script function is another type, it is possible to use MyPy-style type annotations using the types listed above: Example:: @@ -203,6 +208,32 @@ Example:: In our examples, we use comment-based annotations to ensure Python 2 compatibility as well. +2. Annotating Variable Types + +For example, a list by default is assumed to be List[Tensor]. If you would like to +have a list of other types. PyTorch provides annotation functions. + +Example:: + import torch + from torch.jit import Tensor + from typing import List, Tuple + + class ListOfTupleOfTensor(torch.jit.ScriptModule): + def __init__(self): + super(ListOfTupleOfTensor, self).__init__() + + @torch.jit.script_method + def forward(self, x): + # type: (Tensor) -> List[Tuple[Tensor, Tensor]] + + # This annotates the list to be a List[Tuple[Tensor, Tensor]] + returns = torch.jit.annotate(List[Tuple[Tensor, Tensor]], []) + for i in range(10): + returns.append((x, x)) + + return returns + + Expressions ~~~~~~~~~~~