Using torchtyping in vscode, I've found that passing a Tensor to a TorchType generates an error in the type checker:

Tagging the TensorType import with type: ignore as recommended in the FAQ for mypy compatibility doesn't help. Is there any other way to suppress these errors short of tagging every use of a tensor with a tensortype'd sig with type: ignore?

Reproduction

vscode's Pylance language server backs onto the pyright project, and so we can get an easier to examine reproduction by using pyright directly.

Here's a quick script to set up an empty conda env with just torch and torchtyping

mkdir tmp
cd tmp
conda create -p ./.env 
conda activate ./.env
pip install torch==1.9.0 torchtyping==0.1.3

and one more command to install pyright

sudo npm install -g pyright

Then create two files, pyrightconfig.json with contents

{
    "useLibraryCodeForTypes": true,
    "exclude": [".env"]
}

and test.py with contents

import torch
from torchtyping import TensorType

def f(a: TensorType):
    pass

f(torch.zeros())

With that all done, running pyright test.py will give the error:

Loading configuration file at /Users/andy/code/tmp/pyrightconfig.json
Assuming Python version 3.9
Assuming Python platform Darwin
stubPath /Users/andy/code/tmp/typings is not a valid directory.
Searching for source files
Found 1 source file
/Users/andy/code/tmp/test.py
  /Users/andy/code/tmp/test.py:7:3 - error: Argument of type "Tensor" cannot be assigned to parameter "a" of type "TensorType" in function "f"
    "Tensor" is incompatible with "TensorType" (reportGeneralTypeIssues)
1 error, 0 warnings, 0 infos 
Completed in 0.715sec

(Thanks for this awesome work!)

I've often seen applications in which multiple dimensions of the same tensor will be the same size (which size is only known at runtime). It is useful to be able to separately document that those sizes must be the same while acknowledging the differing purposes of these dimensions.

The following example demonstrates the idea---showing three related features that already work in torchtyping as well as a proposed feature:

def func(feats: TensorType["b": ..., "annotator": 3, "word": "words", "feature"],
         predicates: TensorType["b": ..., "annotator": 3, "predicate": "words", "feature"],
         pred_arg_pairs: TensorType["b": ..., "annotator": 3, "predicate": "words", "argument": "words"]):
    # feats has shape (..., 3, words, features)
    # predicates has shape (..., 3, words, features)
    # pred_arg_pairs has shape (..., 3, words, words)
    # the ... b dimensions are checked to be of the same size.

Things that already work:

• dimensions that share names (like "feature") are enforced to share the same dimension size
• named dimensions (like "annotator") can specify a specific dimension size which is enforced
• named ellipses (like "batch") can be used to represent a fixed (but only known at runtime) set of dimensions and corresponding sizes [this is very close to what I want, but (1) I would prefer to not have the extra power of matching an unspecified number of dimensions and (2) as I understand it, ellipses only represent a single variable set of dimensions, but I want to be able to separately constrain multiple sets of dimensions to share respective sizes]

Proposed:

• named dimensions (like "token", "predicate", and "argument") should be able to declare a shared-but-unspecified dimension size given by name ("words" in this example)

Additionally, you would probably want to enforce that, if the specified "size name" matches the name of another dimension (like "word" in the following example), then the sizes of those dimensions should be the same:

def func(feats: TensorType["b": ..., "annotator": 3, "word", "feature"],
         predicates: TensorType["b": ..., "annotator": 3, "predicate": "word", "feature"],
         pred_arg_pairs: TensorType["b": ..., "annotator": 3, "predicate": "word", "argument": "word"]):
    # feats has shape (..., 3, words, features)
    # predicates has shape (..., 3, words, features)
    # pred_arg_pairs has shape (..., 3, words, words)
    # the ... b dimensions are checked to be of the same size.

Thoughts?

We can specify the arbitrary shape of a dimension using the Any type. Example: any_tensor = TensorType[3, 3, Any, torch.float32]

I do not include the tests, to be honest, I am a bit lost on all the tests that you want. Could you clarify them for me @patrick-kidger ?

So that I can do them in the next commit.

Edit: I guess a test/test_any.py and a function test_any_dim in test/test_shape.py

hi! how does the inclusion of variadic generics (pep-646) in python 3.11 affect this project?

will this allow mypy to statically check tensor shapes? maybe a plugin will be required?

just curious! cheers

Hi again :smile: ,

You will find in this pull request some modifications in the code and tests to make them compatible with python 3.7 +. My code is available at the code review.

However, versions 3.7 and 3.8 in the CI are missing (I tested in local on python 3.7 and 3.9).

Hi, I would like to thank you for this cool library. I was desperate not to find a shape typing for pytorch and I had planned to code it myself if it didn't exist.

I think your api is great, however I find that specifying the dimension of any shape to -1 is not very intuitive (I saw that you have many other ways to declare it). One idea is to declare a dimension with any shape using typing.Any. As the library nbtyping does :

from typing import Any
import numpy as np 
from nptyping import NDArray
NDArray[(3, 3, Any), np.float32]

In this case we have typed our array with no constraints on the last dimension. If we apply this modification to your library:

from typing import Any
from torchtyping import TensorType
import torch

TensorType[3, 3, Any, torch.float32]
# Instead of 
TensorType[3, 3, -1, torch.float32]

What do you think of this? If you're interested I can try to make a pull request!

I thank you again for developing this wonderful library.

Since the documentation suggest that mypy integration mostly works, I would have expected the following basic things to work.

import torch
from torchtyping import TensorType  # type: ignore


# I'd expect that mixing TensorType with torch.Tensor works. However in
# strict mode this errors with:
# error: Returning Any from function declared to return "Tensor"
def simple_test_a(x: TensorType[(), float]) -> torch.Tensor:
    return x


# I'd expect that .item() has the type in the annotation, and that e.g.
# annotation the function with `-> str` would be a type error. However
# in strict mode this errors with:
# error: Returning Any from function declared to return "float"
def simple_test_b(x: TensorType[(), float]) -> float:
    return x.item()


# I'd expect that mypy is a aware of actual methods of a tensor and
# gives a type check error when calling garbage.
def simple_test_c(x: TensorType[(), float]) -> None:
    x.asdfasdfasdf()

Note that the latter two work when using the native torch.Tensor as a type annotation.

From what I can tell so far is that mypy actually doesn't understand the meaning of TensorType at all. I thought that "mostly works" means that I can expect mypy to produce meaningful type check errors based on the annotations (operations with illegal shape or dtype combinations etc.).

Am I doing something wrong here, or does "mostly works" that it simply completely ignores the type annotations? Kind of unexpected considering the primary motivation of type annotations is to be used in type checking.

Hello

I'd like to know if there's an easy way to check tensors by name:

import torch
from torch import rand
from torchtyping import TensorType, patch_typeguard, is_named
from typeguard import typechecked

patch_typeguard()  # use before @typechecked


def test():
    t = Test()
    b = t.return_batch()
    o = t.return_other()
    v = t.func(b, o)
    u = t.func(o, b)  # can we have it raise TypeError?


class Test:
    def __init__(self):
        pass

    @typechecked
    def func(
        self,
        x: TensorType["batch"],
        y: TensorType["other"],
    ) -> TensorType["batch", "other"]:
        return torch.outer(x, y)

    def return_batch(self) -> TensorType["batch"]:
        return rand(4)

    def return_other(self) -> TensorType["other"]:
        return rand(3)


test()

Right now, IIUC only dimensions are checked, so in this example there is no error...

I think that I could use is_named in TensorType, but it gets very cumbersome because we also need to use names=... everytime we declare a tensor. This could be OK... but it can get even worse because some pytorch operations don't seem to work with named tensors (outer here! at least with 1.9.1) so we need to rename tensors every 2 lines...

Is is doable to have patch_typeguard(name_check=True), or would it be too complicated to implement? (I think basically I want nominal typing instead of structural typing)

Thanks for your work!

Hi, I really like this project! It's mind blowing for me!!

I have a question. When coding in pycharm using TensorType with a string dimension name, I got a warning about unresolved reference

I am not sure if this is a limitation of pycharm. I am using pycharm 2021.1.1, python 3.7.10 and torchtyping 0.1.4.

Thanks!

Hi! Is it an option to use the torchtyping annotations with other runtime type checkers instead of typeguard?

It is not directly related, but the question came up from the next case. For batch I use dataclass, and typeguard doesn't check in runtime constructor of dataclass. Maybe there is a more clever choice for from what inherit batch?

Consider this function

@typechecked
def mean_squared_error(input: TensorType["batch"], target: TensorType["batch"]):
    d = input - target
    d = d * d
    return torch.mean(d)

The above only allows batches to contain 1-element values (i.e. scalars).

I would like to ensure that the shape of items in the input batch is the same as the shape of items in the target batch, i.e input.shape[1:] == target.shape[1:].

I don't want to hardcode the number of dimensions like for example a batch containing images: input: TensorType["batch", "c", "h", "w"].

Is this currently possible?

As promised, here is the PR to upgrade the library to define a 'torch-like' protocol and use that for the base type rather than using torch.Tensor directly. This lets users perform dimension checking on classes that support a Tensor interface but do not directly inherit from torch.Tensor. I think the change is fairly clear-cut, I have added a test case to demonstrate and verify that dimensions are actually checked. The only question I have is about the change to line 304 in typechecker.py (the last change below). Is this test really necessary? I had to change it to use default construction because protocols don't support isinstance if they have properties.

Hello and thanks for publishing this library! I've really enjoyed reading the design and discussion documents you have posted. However, I am now trying to apply this library in a somewhat broader context. Essentially, I am hoping to use it to improve the linear operator library. linear_operator. The idea of this library is to abstract how tensors are stored to be able to perform matrix operations much more efficiently. I'd really like to use torchtyping to add dimensional and storage type checks to help squash bugs in this code. Unfortunately, torchtyping is configured to run exactly on torch.Tensor objects. My first attempt was just to hack the library to pull out a few class checks. But, doing more reading, I feel like torchtyping could be cleanly improved by using protocols. PEP 544 – Protocols: Structural subtyping (static duck typing). The idea would be to have the library use an abstract tensor protocol rather than tensor directly. This would make the library much more general and I think it could help cleanup the code by making it explicit as to what tensor fields are being used. What do you think / do you have any suggestions on how to add this? @dannyfriar @m4rs-mt

Setup

• pyright: 1.1.263
• pytorch: 1.12.0+cu113
• torchtyping: 0.1.4

Code Example

from torchtyping import TensorType

def example(foo: TensorType["batch"]):
    pass

Problem

Pyright reports the following error: "batch" is not defined

Related issue

The same error is reported by mypy when -1 is omitted: https://github.com/patrick-kidger/torchtyping/issues/35

When I specify a type like TensorType["batch_size", "num_channels", "x", "y"], I get a mypy error like error: Name "batch_size" is not defined for each of the named axes. Is this expected? Am I doing something wrong? This is with the most recent mypy, 0.950.

First off - I've been thinking of almost the same idea as this library for a while because I see runtime errors from tensor shape/dtype mismatches all the time, so glad that there's already something in place!

My initial approach was going to be parsing docstrings of various formats (with an existing library like docstring_parser) and performing validation on these, rather than type annotations. Is that a feature you'd consider accepting into your library? I'd be interested in writing a PR for it with some guidance

For example, I currently write Sphinx style docstrings like this

def forward(self, imgs, tokens):
    """Combine multimodal input features

    :param torch.FloatTensor[N, C, H, W] imgs: Batch of image pixels
        normalized in range of 0-1
    :param torch.LongTensor[N, L] tokens: Vocabularly tokens in sequence
        of length ``L``
    :return torch.FloatTensor[N, C] pred: Predicted probabilities for
        each class
    """

I realize the [N, C, H, W] notation is not quite as rigid as what this project proposes, and that's one reason I've been looking for a more structured approach. But regardless, I do find it nice sometimes to have this information in the docstrings instead of type annotations, particularly for functions with many parameters

n00b to this very cool project, looking to enforce a broadcast-ability pattern where a dimension in one tensor either matches or can be broadcast to (i.e. equals 1) a dimension in another tensor.

@typeguard.typechecked
def mwe(
    x: torchtyping.TensorType[
        ...,
        "foo",
        "bar", # How do we make this "match bar from arg_b or equal 1"?
    ],
    y: torchtyping.TensorType[
        "bar",
    ]) -> torch typing.TensorType[...,"foo","bar"]:
    return x * y

Am I missing an existing way to do this in torchtyping out of the box? Would this need an extension?