vscode/pylance/pyright don't consider a Tensor to be compatible with TensorType

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

[patrick-kidger]

Ach, that’s a shame. Anyway, I’m not sure what more can be done on the torchtyping end, although I’m open to suggestions. Fundamentally I think this is something on pyright’s end, in supporting __class_getitem__.

[patrick-kidger]

Thanks for the thorough repro, that’s really helpful. (Especially as I use neither VSCode nor pyright myself.)

The good news is that I think this can be fixed, and moreover it can be fixed in either of two different ways.

• Option 1 is to switch the annotation from TensorType to TensorType[...], with a literal .... This is equivalent as far as torchtyping is concerned, and this seems to make pyright happy.
• Option 2 is to switch the annotation from TensorType to torch.Tensor, just for those cases where you don’t need the explicit TensorType[...stuff here...]. The explicit TensorType[...stuff here...] should already be fine with pyright.

Does that help?