chore: adding isort to pre-commit

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

.pre-commit-config.yaml

@@ -44,6 +44,7 @@ repos:
     hooks:
     -   id: mypy
         exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools"
+        python_version: "3.11"
   - repo: local
     hooks:
     -   id: dont-commit-upstream

py/torch_tensorrt/_Device.py

@@ -1,17 +1,18 @@
-from typing import Optional, Any, Tuple
 import sys
+from typing import Any, Optional, Tuple
 
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
     from typing_extensions import Self
 
+import warnings
+
 import torch
+from torch_tensorrt import logging
 
 # from torch_tensorrt import _enums
 import tensorrt as trt
-from torch_tensorrt import logging
-import warnings
 
 try:
     from torch_tensorrt import _C

py/torch_tensorrt/_Input.py

@@ -1,8 +1,7 @@
 from enum import Enum
-from typing import List, Dict, Any, Tuple, Optional, Sequence
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import torch
-
 from torch_tensorrt import _enums
 
 

py/torch_tensorrt/__init__.py

@@ -1,16 +1,16 @@
 import ctypes
 import os
-import sys
 import platform
+import sys
+from typing import Dict, List
+
 from packaging import version
 from torch_tensorrt._version import (
     __cuda_version__,
     __cudnn_version__,
     __tensorrt_version__,
 )
 
-from typing import Dict, List
-
 if sys.version_info < (3,):
     raise Exception(
         "Python 2 has reached end-of-life and is not supported by Torch-TensorRT"
@@ -81,11 +81,9 @@ def _find_lib(name: str, paths: List[str]) -> str:
             ctypes.CDLL(_find_lib(lib, LINUX_PATHS))
 
 import torch
-
 from torch_tensorrt._compile import *  # noqa: F403
-from torch_tensorrt._util import *  # noqa: F403
 from torch_tensorrt._enums import *  # noqa: F403
-
+from torch_tensorrt._util import *  # noqa: F403
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from torch_tensorrt import dynamo  # noqa: F401

py/torch_tensorrt/_compile.py

@@ -1,19 +1,16 @@
-from typing import List, Any, Set, Callable, TypeGuard, Optional
+from enum import Enum
+from typing import Any, Callable, List, Optional, Set, TypeGuard
 
+import torch
+import torch.fx
 import torch_tensorrt.ts
-
 from torch_tensorrt import logging
-from torch_tensorrt._Input import Input
 from torch_tensorrt._enums import dtype
-import torch
-import torch.fx
-from enum import Enum
-
-from torch_tensorrt.fx import InputTensorSpec
-from torch_tensorrt.fx.utils import LowerPrecision
-
+from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo.compile import compile as dynamo_compile
+from torch_tensorrt.fx import InputTensorSpec
 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
 
 

py/torch_tensorrt/_enums.py

@@ -0,0 +1,2 @@
+from torch_tensorrt._C import dtype, EngineCapability, TensorFormat  # noqa: F401
+from tensorrt import DeviceType  # noqa: F401

py/torch_tensorrt/_util.py

@@ -1,7 +1,5 @@
-from torch_tensorrt import __version__
-from torch_tensorrt import _C
-
 import torch
+from torch_tensorrt import _C, __version__
 
 
 def dump_build_info() -> None:

py/torch_tensorrt/dynamo/_settings.py

@@ -1,16 +1,17 @@
 from dataclasses import dataclass, field
 from typing import Optional, Set
+
 import torch
 from torch_tensorrt.dynamo._defaults import (
-    PRECISION,
     DEBUG,
-    WORKSPACE_SIZE,
-    MIN_BLOCK_SIZE,
-    PASS_THROUGH_BUILD_FAILURES,
     MAX_AUX_STREAMS,
-    VERSION_COMPATIBLE,
+    MIN_BLOCK_SIZE,
     OPTIMIZATION_LEVEL,
+    PASS_THROUGH_BUILD_FAILURES,
+    PRECISION,
     USE_PYTHON_RUNTIME,
+    VERSION_COMPATIBLE,
+    WORKSPACE_SIZE,
 )
 
 

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -7,7 +7,6 @@
 import torch._dynamo as torchdynamo
 from torch import _guards
 from torch.fx.passes.infra.pass_base import PassResult
-
 from torch_tensorrt.dynamo.utils import req_torch_version
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
     compose_bmm,

py/torch_tensorrt/dynamo/backend/backends.py

@@ -1,25 +1,16 @@
 import logging
-from typing import Sequence, Any, Callable
-import torch
 from functools import partial
-import torch._dynamo as td
+from typing import Any, Callable, Sequence
 
+import torch
+import torch._dynamo as td
+from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
 from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.lowering._decompositions import (
-    get_decompositions,
-)
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import (
-    pre_aot_substitutions,
-)
-from torch_tensorrt.dynamo.lowering._partition import (
-    partition,
-    get_submod_inputs,
-)
-from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 from torch_tensorrt.dynamo.conversion import convert_module
-
-from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
-
+from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
+from torch_tensorrt.dynamo.lowering._partition import get_submod_inputs, partition
+from torch_tensorrt.dynamo.lowering._pre_aot_lowering import pre_aot_substitutions
+from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 
 logger = logging.getLogger(__name__)
 

py/torch_tensorrt/dynamo/compile.py

@@ -1,34 +1,29 @@
-import torch
-import logging
 import collections.abc
-import torch_tensorrt
+import logging
+from typing import Any, List, Optional, Set, Tuple
 
-from typing import Any, Optional, Set, List, Tuple
-from torch_tensorrt import EngineCapability, Device
+import torch
+import torch_tensorrt
 from torch.fx.passes.pass_manager import PassManager
 from torch.fx.passes.splitter_base import SplitResult
-from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
-from torch_tensorrt.dynamo.lowering import (
-    fuse_permute_linear,
-    fuse_permute_matmul,
-)
+from torch_tensorrt import Device, EngineCapability
 from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
-from torch_tensorrt.dynamo.backend.backends import _compile_module
-from torch_tensorrt.dynamo.conversion import convert_module
-
 from torch_tensorrt.dynamo._defaults import (
-    PRECISION,
     DEBUG,
-    WORKSPACE_SIZE,
-    MIN_BLOCK_SIZE,
-    PASS_THROUGH_BUILD_FAILURES,
     MAX_AUX_STREAMS,
-    VERSION_COMPATIBLE,
+    MIN_BLOCK_SIZE,
     OPTIMIZATION_LEVEL,
+    PASS_THROUGH_BUILD_FAILURES,
+    PRECISION,
     USE_PYTHON_RUNTIME,
+    VERSION_COMPATIBLE,
+    WORKSPACE_SIZE,
 )
-
+from torch_tensorrt.dynamo.backend.backends import _compile_module
+from torch_tensorrt.dynamo.conversion import convert_module
+from torch_tensorrt.dynamo.lowering import fuse_permute_linear, fuse_permute_matmul
+from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
+from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
 
 logger = logging.getLogger(__name__)
 

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -1,25 +1,21 @@
 import logging
 import warnings
 from datetime import datetime
-from packaging import version
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set
 
 import numpy
-
-# @manual=//deeplearning/trt/python:py_tensorrt
-import tensorrt as trt
 import torch
 import torch.fx
+from packaging import version
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.shape_prop import TensorMetadata
-
-from torch_tensorrt.fx import CONVERTERS
 from torch_tensorrt import Input
+from torch_tensorrt.fx import CONVERTERS
 from torch_tensorrt.fx.observer import Observer
-from torch_tensorrt.fx.utils import (
-    unified_dtype_converter,
-    Frameworks,
-)
+from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
+
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 

py/torch_tensorrt/dynamo/conversion/conversion.py

@@ -1,11 +1,11 @@
+import io
 from typing import Sequence
+
 import torch
-import io
-from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
-from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
-
+from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 
 import tensorrt as trt
 

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -1,2 +1,2 @@
-from .substitutions import *  # noqa: F403
 from ._fusers import *  # noqa: F403
+from .substitutions import *  # noqa: F403

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -1,7 +1,7 @@
-from typing import Callable, Dict, Any
-import torch
-from torch._decomp import register_decomposition, core_aten_decompositions, OpOverload
+from typing import Any, Callable, Dict
 
+import torch
+from torch._decomp import OpOverload, core_aten_decompositions, register_decomposition
 
 DECOMPOSITIONS: Dict[OpOverload, Callable[..., Any]] = {**core_aten_decompositions()}
 

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -1,18 +1,15 @@
 import logging
-from typing import List, Optional, Sequence, Set, Mapping
+from typing import List, Mapping, Optional, Sequence, Set
 
 import torch
-
-from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
-from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import _get_qualified_name
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupport, SupportDict
-
+from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
 from torch_tensorrt.fx.converter_registry import CONVERTERS
 
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_SINGLE_NODE_PARTITIONS: List[str] = [

py/torch_tensorrt/dynamo/lowering/_pre_aot_lowering.py

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

py/torch_tensorrt/dynamo/lowering/substitutions/__init__.py

@@ -1,2 +1,2 @@
-from .maxpool1d import *  # noqa: F403
 from .einsum import *  # noqa: F403
+from .maxpool1d import *  # noqa: F403

py/torch_tensorrt/dynamo/lowering/substitutions/einsum.py

@@ -1,14 +1,13 @@
-from typing import Dict, Tuple, Any, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence, Tuple
+
 import torch
 from torch._custom_op.impl import custom_op
 from torch.fx.node import Argument, Target
-
+from torch_tensorrt.dynamo.lowering._pre_aot_lowering import register_substitution
 from torch_tensorrt.fx.converter_registry import tensorrt_converter
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import register_substitution
-
 
 @custom_op(
     qualname="tensorrt::einsum",

py/torch_tensorrt/dynamo/lowering/substitutions/maxpool1d.py

@@ -1,15 +1,13 @@
-from typing import Dict, Tuple, Any, Optional
+from typing import Any, Dict, Optional, Tuple
+
 import torch
 from torch._custom_op.impl import custom_op
 from torch.fx.node import Argument, Target
-
+from torch_tensorrt.dynamo.lowering._pre_aot_lowering import register_substitution
 from torch_tensorrt.fx.converter_registry import tensorrt_converter
 from torch_tensorrt.fx.converters import acc_ops_converters
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import register_substitution
-
-
 # This file serves as an example and a tutorial for excluding custom modules from
 # torch.compile tracing. Each required step is labeled with a number indicating the
 # preferable implementation order.

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -1,9 +1,10 @@
-from typing import Any, List, Sequence, Dict, Tuple, Optional
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import torch
+from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
-import torch
-from torch_tensorrt.fx.utils import unified_dtype_converter, Frameworks
 
 
 class PythonTorchTensorRTModule(torch.nn.Module):

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, List, Tuple, Optional
+from typing import Any, List, Optional, Tuple
 
 import torch
 from torch_tensorrt._Device import Device

py/torch_tensorrt/dynamo/runtime/__init__.py

@@ -0,0 +1,2 @@
+from ._PythonTorchTensorRTModule import PythonTorchTensorRTModule # noqa: F401
+from ._TorchTensorRTModule import TorchTensorRTModule  # noqa: F401