Fix TensorType to return new TensorTypes for shapes instead of TensorType instances

[abock]

This ensures that TensorType is never actually instantiated and instead new TensorType subclasses are created (and cached) for each shape.

This is a stepping stone to a fully generic TensorType where the dtype and shape are modeled into the type itself, but it appears that won't be easily realized until we can depend on Python 3.11.

Importantly, this adds unit tests to ensure if we can move to a proper generic that behavior will remain the same.

Fixes #219

[github-advanced-security]

lintrunner found more than 10 potential problems in the proposed changes. Check the Files changed tab for more details.

[codecov]

Codecov Report

Merging #228 (fcd664e) into main (ed41071) will increase coverage by 0.08%.
The diff coverage is 98.83%.

@@            Coverage Diff             @@
##             main     #228      +/-   ##
==========================================
+ Coverage   75.57%   75.66%   +0.08%     
==========================================
  Files          89       90       +1     
  Lines        7248     7274      +26     
==========================================
+ Hits         5478     5504      +26     
  Misses       1770     1770              

Impacted Files	Coverage Δ
onnxscript/test/onnx_types_test.py	97.22% <97.22%> (ø)
onnxscript/onnx_types.py	96.15% <100.00%> (+0.53%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[justinchuby]

LGTM🚀

[justinchuby]

I would import abc and use abc.ABC

[abock]

Any particular reason?

[justinchuby]

Mostly because of this https://github.com/microsoft/onnx-script/wiki/Coding-style#avoid-relative-imports-import-only-modules, as suggested by https://google.github.io/styleguide/pyguide.html#23-packages.

ABC is usually unique enough, but I think maintaining a single rule of importing only modules can save us from deciding what to import as modules and what not to.

[gramalingam]

LGTM, thanks. One thing I am not entirely clear is how well other tools, like mypy, will be when dealing with onnxscript code with type annotations, but that's a different question, may be worth talking about in our meeting. E.g., would using something like Annotated[TensorType, Dtype(dtype), Shape(shape)] work better with other type-checkers and tools?

[abock]

As for tooling, I am very unimpressed with mypy in particular. There is an actual bug in mypy that is causing it to not see the __class_getitem__ and then complain that the type is not generic (it is not) or subscriptable (it is).

Python 3.11 provides support for variadic generics and syntax sugar to promote Python literals to Literal[T]. We could make it fully generic today via TypeVarTuple and Unpack through typing_extensions, but the usage syntax is gross.

Today:

def F(tensor: FLOAT[Literal[128], Literal["M"], Literal[1024]]): ...

def G(tensor: FLOAT[Literal[Ellipsis]]): ...

With 3.11:

def F(tensor: FLOAT[128, "M", 1024]): ...

def G(tensor: FLOAT[...]): ...