Typing in Pytorch

sh0416: <p>Hi,</p>
<p>I use type annotation for my code.</p>
<p>Is there any typing annotation guideline for pytorch?</p>
<p>I want to do something like this.</p>
<pre><code class="lang-auto">class MyModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.linear = nn.Linear(10, 4)

  def forward(self, x: torch.Tensor[torch.float, [B, 10]]) -&gt; torch.Tensor[torch.float, [B, 4]]:
    return self.linear(x)
</code></pre>
<p>I guess torchScript or some other approach, e.g. View(?), could support this one, but I couldn’t find official or near-official document.</p>
<p>Is there anyone who know about this, help me <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>
<p>Thanks,</p>

davince: <p>Hi <a class="mention" href="/u/sh0416">@sh0416</a>  here is a sample of code i cameby while browsing <a href="https://guide.allennlp.org/training-and-prediction#1" rel="noopener nofollow ugc">Allen Nlp</a> guides and source code. I hope this points you to the right direction. Sorry I just pasted there example codes</p>
<pre><code class="lang-python">from typing import Dict, Iterable, List

from allennlp.data import DatasetReader, Instance
from allennlp.data.fields import LabelField, TextField
from allennlp.data.token_indexers import TokenIndexer, SingleIdTokenIndexer
from allennlp.data.tokenizers import Token, Tokenizer, WhitespaceTokenizer
</code></pre>
<pre><code class="lang-python">class ClassificationTsvReader(DatasetReader):
    def __init__(self,
                 lazy: bool = False,
                 tokenizer: Tokenizer = None,
                 token_indexers: Dict[str, TokenIndexer] = None,
                 max_tokens: int = None):
        super().__init__(lazy)
        self.tokenizer = tokenizer or WhitespaceTokenizer()
        self.token_indexers = token_indexers or {'tokens': SingleIdTokenIndexer()}
        self.max_tokens = max_tokens

    def _read(self, file_path: str) -&gt; Iterable[Instance]:
        with open(file_path, 'r') as lines:
            for line in lines:
                text, sentiment = line.strip().split('\t')
                tokens = self.tokenizer.tokenize(text)
                if self.max_tokens:
                    tokens = tokens[:self.max_tokens]
                text_field = TextField(tokens, self.token_indexers)
                label_field = LabelField(sentiment)
                fields = {'text': text_field, 'label': label_field}
                yield Instance(fields)

dataset_reader = ClassificationTsvReader(max_tokens=64)
instances = dataset_reader.read("quick_start/data/movie_review/train.tsv")

for instance in instances[:10]:
    print(instance)
</code></pre>
<p>or this</p>
<pre><code class="lang-python">class SimpleClassifier(Model):
    def __init__(self,
                 vocab: Vocabulary,
                 embedder: TextFieldEmbedder,
                 encoder: Seq2VecEncoder):
        super().__init__(vocab)
        self.embedder = embedder
        self.encoder = encoder
        num_labels = vocab.get_vocab_size("labels")
        self.classifier = torch.nn.Linear(encoder.get_output_dim(), num_labels)

    def forward(self,
                text: Dict[str, torch.Tensor],
                label: torch.Tensor) -&gt; Dict[str, torch.Tensor]:
        # Shape: (batch_size, num_tokens, embedding_dim)
        embedded_text = self.embedder(text)
        # Shape: (batch_size, num_tokens)
        mask = util.get_text_field_mask(text)
        # Shape: (batch_size, encoding_dim)
        encoded_text = self.encoder(embedded_text, mask)
        # Shape: (batch_size, num_labels)
        logits = self.classifier(encoded_text)
        # Shape: (batch_size, num_labels)
        probs = torch.nn.functional.softmax(logits, dim=-1)
        # Shape: (1,)
        loss = torch.nn.functional.cross_entropy(logits, label)
        return {'loss': loss, 'probs': probs}
</code></pre>

sh0416: <p>Thanks for the suggestion!</p>
<p>The second one is closer to the one I want. (Check the shape of tensor)</p>

kkakkavas: <p>You could try <a href="https://github.com/patrick-kidger/torchtyping/tree/master" rel="noopener nofollow ugc">torchtyping</a> <code>pip install torchtyping</code></p>
<p>Or jaxtyping that also support torch and numpy.</p>
<p>a simple example is given below from the github readme</p>
<p>Turn this</p>
<pre data-code-wrap="python"><code class="lang-python">def batch_outer_product(x: torch.Tensor, y: torch.Tensor) -&gt; torch.Tensor:
    # x has shape (batch, x_channels)
    # y has shape (batch, y_channels)
    # return has shape (batch, x_channels, y_channels)

    return x.unsqueeze(-1) * y.unsqueeze(-2)
</code></pre>
<p>To this</p>
<pre data-code-wrap="python"><code class="lang-python">def batch_outer_product(x:   TensorType["batch", "x_channels"],
                        y:   TensorType["batch", "y_channels"]
                        ) -&gt; TensorType["batch", "x_channels", "y_channels"]:

    return x.unsqueeze(-1) * y.unsqueeze(-2)
</code></pre>