chore: ready to start review

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

Author: narendasan

.pre-commit-config.yaml

@@ -5,6 +5,7 @@ repos:
     hooks:
       - id: check-yaml
       - id: trailing-whitespace
+        exclude: ^docs
       - id: check-added-large-files
         args:
           - --maxkb=1000
@@ -13,11 +14,7 @@ repos:
       - id: mixed-line-ending
         args:
           - --fix=lf
-  - repo: https://github.com/psf/black
-    rev: 23.7.0
-    hooks:
-      - id: black
-        exclude: ^examples/custom_converters/elu_converter/setup.py
+        exclude: ^docs
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v16.0.6
     hooks:
@@ -30,21 +27,30 @@ repos:
         args:
           - --warnings=all
     -   id: buildifier-lint
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    # Ruff version.
-    rev: v0.0.278
-    hooks:
-      - id: ruff
   - repo: https://github.com/abravalheri/validate-pyproject
     rev: v0.13
     hooks:
       - id: validate-pyproject
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        name: isort (python)
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: 'v1.4.1'
     hooks:
     -   id: mypy
-        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools"
-        python_version: "3.11"
+        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools|^docs|noxfile.py|setup.py|versions.py"
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.0.278
+    hooks:
+      - id: ruff
+  - repo: https://github.com/psf/black
+    rev: 23.7.0
+    hooks:
+      - id: black
+        exclude: ^examples/custom_converters/elu_converter/setup.py|^docs
   - repo: local
     hooks:
     -   id: dont-commit-upstream

core/conversion/evaluators/aten.cpp

@@ -104,7 +104,12 @@ DEFINE_ARITHMATIC_TWO_INPUT_EVALUATOR(
         "aten::pow.float_int(float a, int b) -> (float)",
     }));
 
-DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(and, "aten::__and__", a&& b, bool, std::set<std::string>({"aten::__and__(int a, int b) -> (bool)", "aten::__and__.bool(bool a, bool b) -> (bool)"}));
+DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
+    and,
+    "aten::__and__",
+    a&& b,
+    bool,
+    std::set<std::string>({"aten::__and__(int a, int b) -> (bool)", "aten::__and__.bool(bool a, bool b) -> (bool)"}));
 DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(or, "aten::__or__", a || b, bool, {"aten::__or__(int a, int b) -> (bool)"});
 DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
     xor,

docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py

@@ -16,6 +16,7 @@
 
 # %%
 
+
 # We begin by defining a model
 class Model(torch.nn.Module):
     def __init__(self) -> None:

docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py

@@ -16,6 +16,7 @@
 
 # %%
 
+
 # We begin by defining a model
 class Model(torch.nn.Module):
     def __init__(self) -> None:

py/torch_tensorrt/__init__.py

@@ -4,13 +4,14 @@
 import sys
 from typing import Dict, List
 
-from packaging import version
 from torch_tensorrt._version import (
     __cuda_version__,
     __cudnn_version__,
     __tensorrt_version__,
 )
 
+from packaging import version
+
 if sys.version_info < (3,):
     raise Exception(
         "Python 2 has reached end-of-life and is not supported by Torch-TensorRT"
@@ -82,10 +83,10 @@ def _find_lib(name: str, paths: List[str]) -> str:
 
 import torch
 from torch_tensorrt._compile import *  # noqa: F403
+from torch_tensorrt._Device import Device  # noqa: F401
 from torch_tensorrt._enums import *  # noqa: F403
-from torch_tensorrt._util import *  # noqa: F403
-from torch_tensorrt._Input import Input # noqa: F401
-from torch_tensorrt._Device import Device # noqa: F401
+from torch_tensorrt._Input import Input  # noqa: F401
+from torch_tensorrt._utils 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,5 +1,5 @@
 from enum import Enum
-from typing import Any, Callable, List, Optional, Set, TypeGuard
+from typing import Any, Callable, List, Optional, Sequence, Set, TypeGuard
 
 import torch
 import torch.fx
@@ -15,13 +15,13 @@
 
 
 def _non_fx_input_interface(
-    inputs: List[Input | torch.Tensor | InputTensorSpec],
+    inputs: Sequence[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[Input | torch.Tensor]]:
     return all(isinstance(i, torch.Tensor | Input) for i in inputs)
 
 
 def _fx_input_interface(
-    inputs: List[Input | torch.Tensor | InputTensorSpec],
+    inputs: Sequence[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[InputTensorSpec | torch.Tensor]]:
     return all(isinstance(i, torch.Tensor | InputTensorSpec) for i in inputs)
 
@@ -97,7 +97,7 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 def compile(
     module: Any,
     ir: str = "default",
-    inputs: Optional[List[Input | torch.Tensor | InputTensorSpec]] = None,
+    inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None,
     enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,
 ) -> (
@@ -201,7 +201,7 @@ def compile(
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 
 
-def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Callable[..., Any]:
+def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Any:
     """
     Returns a boxed model which is the output of torch.compile.
     This does not compile the model to TRT. Execute this model on
@@ -216,8 +216,8 @@ def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Callable[..., Any]:
 
 def convert_method_to_trt_engine(
     module: Any,
-    inputs: List[Input | torch.Tensor],
-    method_name: str,
+    method_name: str = "forward",
+    inputs: Optional[Sequence[Input | torch.Tensor]] = None,
     ir: str = "default",
     enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,
@@ -266,7 +266,7 @@ def convert_method_to_trt_engine(
                 "Module was provided as a torch.nn.Module, trying to script the module with torch.jit.script. In the event of a failure please preconvert your module to TorchScript",
             )
             ts_mod = torch.jit.script(module)
-        return torch_tensorrt.ts.convert_method_to_trt_engine(
+        return torch_tensorrt.ts.convert_method_to_trt_engine(  # type: ignore[no-any-return]
             ts_mod,
             inputs=inputs,
             method_name=method_name,

py/torch_tensorrt/_enums.py

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

py/torch_tensorrt/_util.py renamed to py/torch_tensorrt/_utils.py

@@ -1,5 +1,8 @@
+from typing import Any
+
 import torch
-from torch_tensorrt import _C, __version__
+from torch_tensorrt import _C
+from torch_tensorrt._version import __version__
 
 
 def dump_build_info() -> None:
@@ -30,7 +33,7 @@ def set_device(gpu_id: int) -> None:
     _C.set_device(gpu_id)
 
 
-def sanitized_torch_version() -> str:
+def sanitized_torch_version() -> Any:
     return (
         torch.__version__
         if ".nv" not in torch.__version__

py/torch_tensorrt/dynamo/_SourceIR.py

@@ -9,7 +9,7 @@ class SourceIR(Enum):
     TORCHTRT_LOWERED = auto()
     UNKNOWN = auto()
 
-    def __str__(self):
+    def __str__(self) -> str:
         if self == SourceIR.NN:
             return "nn"
         elif self == SourceIR.ACC:

py/torch_tensorrt/dynamo/__init__.py

@@ -1,13 +1,12 @@
+from torch_tensorrt._utils import sanitized_torch_version
+
 from packaging import version
-from torch_tensorrt._util import sanitized_torch_version
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from ._settings import *  # noqa: F403
-    from .conversion import * # noqa: F403
-    from .aten_tracer import trace # noqa: F403
-    from .conversion.converter_registry import (
-        DYNAMO_CONVERTERS, # noqa: F403
-        dynamo_tensorrt_converter, # noqa: F403
-    )
-    from .compile import compile # noqa: F403
-    from ._SourceIR import SourceIR # noqa: F403
+    from ._SourceIR import SourceIR  # noqa: F403
+    from .aten_tracer import trace  # noqa: F403
+    from .compile import compile  # noqa: F403
+    from .conversion import *  # noqa: F403
+    from .conversion.converter_registry import DYNAMO_CONVERTERS  # noqa: F403
+    from .conversion.converter_registry import dynamo_tensorrt_converter  # noqa: F403

py/torch_tensorrt/dynamo/_settings.py

@@ -9,13 +9,10 @@
     OPTIMIZATION_LEVEL,
     PASS_THROUGH_BUILD_FAILURES,
     PRECISION,
-    USE_PYTHON_RUNTIME,
-<<<<<<< HEAD
     TRUNCATE_LONG_AND_DOUBLE,
-=======
+    USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
->>>>>>> e39abb60d (chore: adding isort to pre-commit)
 )
 
 

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -1,11 +1,10 @@
 import copy
 import sys
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
 
 import torch
 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 (
@@ -23,11 +22,7 @@
 )
 from typing_extensions import TypeAlias
 
-Value: TypeAlias = Union[
-    Tuple["Value", ...],
-    List["Value"],
-    Dict[str, "Value"],
-]
+Value: TypeAlias = Tuple["Value", ...] | List["Value"] | Dict[str, "Value"]
 
 
 class DynamoConfig:
@@ -96,7 +91,7 @@ def dynamo_trace(
     aten_graph: bool,
     tracing_mode: str = "real",
     dynamo_config: Optional[DynamoConfig] = None,
-) -> Tuple[torch.fx.GraphModule, Set[_guards.Guard]]:
+) -> Any:  # Tuple[torch.fx.GraphModule, Set[_guards.Guard]]:
     """
     TODO: Once we fully migrate to torchdynamo frontend, we will remove
     this config option alltogether.  For now, it helps with quick
@@ -150,11 +145,11 @@ def trace(
     fx_module, __package__ = dynamo_trace(model, inputs, True, "symbolic")
     print(fx_module.graph)
     for passes in passes_list:
-        pr: PassResult = passes(fx_module)  # type: ignore[assignment] #The type hints in fx are wrong
+        pr: PassResult = passes(fx_module)
         fx_module = pr.graph_module
 
     fx_module(*inputs)
 
     fx_module = run_const_fold(fx_module)
     print(fx_module.graph)
-    return fx_module  # type: ignore[no-any-return]
+    return fx_module

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -1 +1 @@
-from .backends import torch_tensorrt_backend # noqa: F401
+from .backends import torch_tensorrt_backend  # noqa: F401

py/torch_tensorrt/dynamo/backend/backends.py

@@ -4,30 +4,12 @@
 
 import torch
 import torch._dynamo as td
-<<<<<<< HEAD
-
+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,
     repair_long_or_double_inputs,
 )
-
-=======
->>>>>>> e39abb60d (chore: adding isort to pre-commit)
-from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
-from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.conversion import convert_module
 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
@@ -64,7 +46,7 @@ def aot_torch_tensorrt_aten_backend(
     return aot_module_simplified(
         gm,
         sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend),  # type: ignore[no-untyped-call]
+        fw_compiler=make_boxed_compiler(custom_backend),
         decompositions=get_decompositions(),
     )