Patch for tensor types

[justinchuby]

Now mypy can correctly check types:

It still doesn't like that that they are not generics, but I think we can disable type-arg and name-defined for now.

  Error (MYPY) type-arg
    "FLOAT" expects no type arguments, but 1 given

        28  |        self.assertSameGraph(static_shape, static_shape_txt)
        29  |
        30  |        @script()
    >>> 31  |        def symbolic_shape(A: FLOAT["N"], B: FLOAT["N"]) -> FLOAT["N"]:  # noqa: F821
        32  |            C = op.Add(A, B)
        33  |            return C
        34  |

  Error (MYPY) name-defined
    Name "N" is not defined

        28  |        self.assertSameGraph(static_shape, static_shape_txt)
        29  |
        30  |        @script()
    >>> 31  |        def symbolic_shape(A: FLOAT["N"], B: FLOAT["N"]) -> FLOAT["N"]:  # noqa: F821
        32  |            C = op.Add(A, B)
        33  |            return C
        34  |

[codecov]

Codecov Report

Merging #231 (240f226) into abock/uninstantiated-tensor-types (08dc10d) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

@@                          Coverage Diff                          @@
##           abock/uninstantiated-tensor-types     #231      +/-   ##
=====================================================================
- Coverage                              75.62%   75.61%   -0.02%     
=====================================================================
  Files                                     90       90              
  Lines                                   7241     7238       -3     
=====================================================================
- Hits                                    5476     5473       -3     
  Misses                                  1765     1765              

Impacted Files	Coverage Δ
onnxscript/irbuilder.py	78.96% <100.00%> (ø)
onnxscript/onnx_types.py	95.53% <100.00%> (-0.62%)	⬇️
onnxscript/test/onnx_types_test.py	95.83% <100.00%> (-1.31%)	⬇️
onnxscript/test/type_annotation_test.py	74.28% <100.00%> (ø)
onnxscript/test/converter_test.py	87.55% <0.00%> (+0.23%)	⬆️

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

[justinchuby]

I don't think this works because FLOAT is now an instance of TensorType, not a subclass.

[justinchuby]

Will wait for

• mypy does not handle __class_getitem__ python/mypy#11501

[justinchuby]

Alternatively, we will use this and run mypy with python3.11. So mypy is happy and we can pass in literals as types even with python <3.11.

Will need to fix dtype and shape.

from typing import Any, Generic, TypeVar, Union, Tuple, Literal
import typing
from typing_extensions import TypeVarTuple, Unpack
DimType = Union[int, str, type(None)]
ShapeType = Union[Tuple[DimType, ...], DimType, type(Ellipsis)]
TDType = TypeVar("TDType")
TShape = TypeVarTuple("TShape")
class _DTYPE_FLOAT:
    dtype = 100
class TensorType(Generic[TDType, Unpack[TShape]]):
    def __class_getitem__(cls, *items: Any):
        wrapped = []
        for item in items[0]:
            if type(item) in {int, str, type(...), type(None)}:
                wrapped.append(Literal[item])
            else:
                wrapped.append(item)
        return super().__class_getitem__(
            tuple(wrapped)
        )
    def dtype(self) -> TDType:
        return typing.get_args(self.__orig_class__)[0]
    def shape(self) -> TShape:
        return typing.get_args(self.__orig_class__)[1:]
class FLOAT(TensorType[_DTYPE_FLOAT, Unpack[TShape]]):
    pass
print(TensorType[int, 1, 2])
print(TensorType[int, 1, 2]().dtype())
print(TensorType[int, 1, 2]().shape())
M = TypeVar("M", bound=int)
print(FLOAT[1, 2, ..., "N", M])