chore: ready to start review

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

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,26 @@ 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
+        python_version: "3.11"
   - 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

@@ -103,7 +103,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

@@ -82,10 +82,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, Set, TypeGuard, Sequence
 
 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,

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 py/torch_tensorrt/_utils.pypy/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._version import __version__
+from torch_tensorrt import _C
 
 
 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,11 @@
 from packaging import version
-from torch_tensorrt._util import sanitized_torch_version
+from torch_tensorrt._utils 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