Merge pull request #2189 from pytorch/py38_compatibility

Py38 compatibility

Author: narendasan

py/torch_tensorrt/_Input.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from enum import Enum
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
@@ -32,11 +34,11 @@ class _ShapeMode(Enum):
     shape: Optional[
         Tuple[int, ...] | Dict[str, Tuple[int, ...]]
     ] = None  #: Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
-    dtype: _enums.dtype = (  # type: ignore[name-defined]
+    dtype: _enums.dtype = (
         _enums.dtype.unknown
     )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
     _explicit_set_dtype: bool = False
-    format: _enums.TensorFormat = (  # type: ignore[name-defined]
+    format: _enums.TensorFormat = (
         _enums.TensorFormat.contiguous
     )  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
@@ -208,7 +210,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
             return False
 
     @staticmethod
-    def _parse_dtype(dtype: Any) -> _enums.dtype:  # type: ignore[name-defined]
+    def _parse_dtype(dtype: Any) -> _enums.dtype:
         if isinstance(dtype, torch.dtype):
             if dtype == torch.long:
                 return _enums.dtype.long
@@ -236,7 +238,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:  # type: ignore[name-defined]
             )
 
     @staticmethod
-    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:  # type: ignore[name-defined]
+    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:
         if dtype == _enums.dtype.long:
             return torch.long
         elif dtype == _enums.dtype.int32:
@@ -255,7 +257,7 @@ def is_trt_dtype(self) -> bool:
         return bool(self.dtype != _enums.dtype.long)
 
     @staticmethod
-    def _parse_format(format: Any) -> _enums.TensorFormat:  # type: ignore[name-defined]
+    def _parse_format(format: Any) -> _enums.TensorFormat:
         if isinstance(format, torch.memory_format):
             if format == torch.contiguous_format:
                 return _enums.TensorFormat.contiguous

py/torch_tensorrt/_compile.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from enum import Enum
-from typing import Any, Callable, List, Optional, Sequence, Set, TypeGuard
+from typing import Any, Callable, List, Optional, Sequence, Set
 
 import torch
 import torch.fx
@@ -12,6 +14,7 @@
 from torch_tensorrt.fx.lower import compile as fx_compile
 from torch_tensorrt.fx.utils import LowerPrecision
 from torch_tensorrt.ts._compiler import compile as torchscript_compile
+from typing_extensions import TypeGuard
 
 
 def _non_fx_input_interface(

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import copy
 import sys
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 import torch._dynamo as torchdynamo
@@ -22,7 +24,7 @@
 )
 from typing_extensions import TypeAlias
 
-Value: TypeAlias = Tuple["Value", ...] | List["Value"] | Dict[str, "Value"]
+Value: TypeAlias = Union[Tuple["Value", ...], List["Value"], Dict[str, "Value"]]
 
 
 class DynamoConfig:

py/torch_tensorrt/dynamo/backend/backends.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from functools import partial
 from typing import Any, Callable, Sequence

py/torch_tensorrt/dynamo/compile.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import collections.abc
 import logging
 from typing import Any, List, Optional, Set, Tuple

py/torch_tensorrt/dynamo/conversion/conversion.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 from typing import Sequence
 

py/torch_tensorrt/dynamo/conversion/converter_registry.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from dataclasses import dataclass, field
 from enum import Enum, auto
@@ -28,7 +30,7 @@
         Dict[str, Argument],
         str,
     ],
-    TRTTensor | Sequence[TRTTensor],
+    Union[TRTTensor, Sequence[TRTTensor]],
 ]
 
 

py/torch_tensorrt/dynamo/conversion/impl/shape.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import List, Optional, Tuple
 
 import numpy as np

py/torch_tensorrt/dynamo/conversion/truncate_long_and_double.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Optional, Sequence, Set
 
 import torch

py/torch_tensorrt/dynamo/lowering/_pre_aot_lowering.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 import logging
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, Optional, Type, TypeAlias
+from typing import Any, Callable, Dict, Optional, Type
 
 import torch
 from torch._ops import OpOverload
 from torch.fx import GraphModule, Node
+from typing_extensions import TypeAlias
 
 logger = logging.getLogger(__name__)