Pyright reports an error with named axis

[Luceurre]

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: #35

[patrick-kidger]

This is expected. Static type checking is (unfortunately) fundamentally incompatible with annotating arrays.

You should either add the appropriate flag to disable pyright's checking here, or define batch = None elsewhere in the file so that pyright thinks this is a forward reference.

[wookayin]

Can we use something like TypeVar (instead of string literal for the type annotation) instead?

e.g.

Batch = typing.TypeVar

def example(foo: TensorType[Batch]):
    pass

or

Batch = torchtyping.AxisVar("Batch")   # a hypothetical API

def example(foo: TensorType[Batch]):
    pass

UPDATE: seems relevant to #37 (PEP-646)