chore: All code is now flake8 compliant

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

Author: narendasan

.pre-commit-config.yaml

@@ -30,6 +30,11 @@ 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:
@@ -47,4 +52,3 @@ repos:
         exclude: "^.pre-commit-config.yaml"
         language: pygrep
         types: [text]
-

py/torch_tensorrt/_Device.py

@@ -1,4 +1,4 @@
-from typing import TypeVar, Optional, Any, Tuple
+from typing import Optional, Any, Tuple
 import sys
 
 if sys.version_info >= (3, 11):
@@ -15,7 +15,7 @@
 
 try:
     from torch_tensorrt import _C
-except:
+except ImportError:
     warnings.warn(
         "Unable to import torchscript frontend core and torch-tensorrt runtime. Some dependent features may be unavailable."
     )

py/torch_tensorrt/_Input.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Dict, Any, Tuple, Optional, Union, Sequence
+from typing import List, Dict, Any, Tuple, Optional, Sequence
 
 import torch
 
@@ -86,7 +86,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             self.shape_mode = Input._ShapeMode.STATIC
 
         elif len(args) == 0:
-            if not ("shape" in kwargs) and not (
+            if "shape" not in kwargs and not (
                 all(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"])
             ):
                 raise ValueError(
@@ -287,8 +287,8 @@ def _parse_tensor_domain(
             domain_lo, domain_hi = domain
 
             # Validate type and provided values for domain
-            valid_type_lo = isinstance(domain_lo, int) or isinstance(domain_lo, float)
-            valid_type_hi = isinstance(domain_hi, int) or isinstance(domain_hi, float)
+            valid_type_lo = isinstance(domain_lo, (int, float))
+            valid_type_hi = isinstance(domain_hi, (int, float))
 
             if not valid_type_lo:
                 raise ValueError(
@@ -394,12 +394,10 @@ def example_tensor(
             if optimization_profile_field is not None:
                 try:
                     assert any(
-                        [
-                            optimization_profile_field == field_name
-                            for field_name in ["min_shape", "opt_shape", "max_shape"]
-                        ]
+                        optimization_profile_field == field_name
+                        for field_name in ["min_shape", "opt_shape", "max_shape"]
                     )
-                except:
+                except AssertionError:
                     raise ValueError(
                         "Invalid field name, expected one of min_shape, opt_shape, max_shape"
                     )

py/torch_tensorrt/__init__.py

@@ -1,12 +1,9 @@
 import ctypes
-import glob
 import os
 import sys
 import platform
-import warnings
 from packaging import version
 from torch_tensorrt._version import (
-    __version__,
     __cuda_version__,
     __cudnn_version__,
     __tensorrt_version__,
@@ -38,8 +35,8 @@ def _find_lib(name: str, paths: List[str]) -> str:
 
 
 try:
-    import tensorrt
-except:
+    import tensorrt  # noqa: F401
+except ImportError:
     cuda_version = _parse_semver(__cuda_version__)
     cudnn_version = _parse_semver(__cudnn_version__)
     tensorrt_version = _parse_semver(__tensorrt_version__)
@@ -85,20 +82,14 @@ def _find_lib(name: str, paths: List[str]) -> str:
 
 import torch
 
-from torch_tensorrt._compile import *
-from torch_tensorrt._util import *
-from torch_tensorrt import ts
-from torch_tensorrt import ptq
-from torch_tensorrt._enums import *
-from torch_tensorrt import logging
-from torch_tensorrt._Input import Input
-from torch_tensorrt._Device import Device
+from torch_tensorrt._compile import *  # noqa: F403
+from torch_tensorrt._util import *  # noqa: F403
+from torch_tensorrt._enums import *  # noqa: F403
 
-from torch_tensorrt import fx
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
-    from torch_tensorrt import dynamo
-    from torch_tensorrt.dynamo import backend
+    from torch_tensorrt import dynamo  # noqa: F401
+    from torch_tensorrt.dynamo import backend  # noqa: F401
 
 
 def _register_with_torch() -> None:

py/torch_tensorrt/_compile.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any, Set, Union, Callable, TypeGuard
+from typing import List, Any, Set, Callable, TypeGuard, Optional
 
 import torch_tensorrt.ts
 
@@ -9,21 +9,24 @@
 import torch.fx
 from enum import Enum
 
-import torch_tensorrt.fx
 from torch_tensorrt.fx import InputTensorSpec
 from torch_tensorrt.fx.utils import LowerPrecision
 
+from torch_tensorrt.dynamo.compile import compile as dynamo_compile
+from torch_tensorrt.fx.lower import compile as fx_compile
+from torch_tensorrt.ts._compiler import compile as torchscript_compile
+
 
 def _non_fx_input_interface(
     inputs: List[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[Input | torch.Tensor]]:
-    return all([isinstance(i, torch.Tensor | Input) for i in inputs])
+    return all(isinstance(i, torch.Tensor | Input) for i in inputs)
 
 
 def _fx_input_interface(
     inputs: List[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[InputTensorSpec | torch.Tensor]]:
-    return all([isinstance(i, torch.Tensor | InputTensorSpec) for i in inputs])
+    return all(isinstance(i, torch.Tensor | InputTensorSpec) for i in inputs)
 
 
 class _IRType(Enum):
@@ -58,10 +61,10 @@ def _parse_module_type(module: Any) -> _ModuleType:
 
 
 def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
-    module_is_tsable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.ts]])
-    module_is_fxable = any([module_type == t for t in [_ModuleType.nn, _ModuleType.fx]])
+    module_is_tsable = any(module_type == t for t in [_ModuleType.nn, _ModuleType.ts])
+    module_is_fxable = any(module_type == t for t in [_ModuleType.nn, _ModuleType.fx])
 
-    ir_targets_torchscript = any([ir == opt for opt in ["torchscript", "ts"]])
+    ir_targets_torchscript = any(ir == opt for opt in ["torchscript", "ts"])
     ir_targets_fx = ir == "fx"
     ir_targets_dynamo = ir == "dynamo"
     ir_targets_torch_compile = ir == "torch_compile"
@@ -97,8 +100,8 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 def compile(
     module: Any,
     ir: str = "default",
-    inputs: List[Input | torch.Tensor | InputTensorSpec] = [],
-    enabled_precisions: Set[torch.dtype | dtype] = set([torch.float]),
+    inputs: Optional[List[Input | torch.Tensor | InputTensorSpec]] = None,
+    enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,
 ) -> (
     torch.nn.Module | torch.jit.ScriptModule | torch.fx.GraphModule | Callable[..., Any]
@@ -138,6 +141,11 @@ def compile(
     Returns:
         torch.nn.Module: Compiled Module, when run it will execute via TensorRT
     """
+    input_list = inputs if inputs is not None else []
+    enabled_precisions_set = (
+        enabled_precisions if enabled_precisions is not None else {torch.float}
+    )
+
     module_type = _parse_module_type(module)
     target_ir = _get_target_ir(module_type, ir)
     if target_ir == _IRType.ts:
@@ -148,45 +156,50 @@ def compile(
                 "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)
-        assert _non_fx_input_interface(inputs)
-        compiled_ts_module: torch.jit.ScriptModule = torch_tensorrt.ts.compile(
-            ts_mod, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
+        assert _non_fx_input_interface(input_list)
+        compiled_ts_module: torch.jit.ScriptModule = torchscript_compile(
+            ts_mod,
+            inputs=input_list,
+            enabled_precisions=enabled_precisions_set,
+            **kwargs,
         )
         return compiled_ts_module
     elif target_ir == _IRType.fx:
         if (
-            torch.float16 in enabled_precisions
-            or torch_tensorrt.dtype.half in enabled_precisions
+            torch.float16 in enabled_precisions_set
+            or torch_tensorrt.dtype.half in enabled_precisions_set
         ):
             lower_precision = LowerPrecision.FP16
         elif (
-            torch.float32 in enabled_precisions
-            or torch_tensorrt.dtype.float in enabled_precisions
+            torch.float32 in enabled_precisions_set
+            or torch_tensorrt.dtype.float in enabled_precisions_set
         ):
             lower_precision = LowerPrecision.FP32
         else:
-            raise ValueError(f"Precision {enabled_precisions} not supported on FX")
+            raise ValueError(f"Precision {enabled_precisions_set} not supported on FX")
 
-        assert _fx_input_interface(inputs)
-        compiled_fx_module: torch.nn.Module = torch_tensorrt.fx.compile(
+        assert _fx_input_interface(input_list)
+        compiled_fx_module: torch.nn.Module = fx_compile(
             module,
-            inputs,
+            input_list,
             lower_precision=lower_precision,
             explicit_batch_dimension=True,
             dynamic_batch=False,
             **kwargs,
         )
         return compiled_fx_module
     elif target_ir == _IRType.dynamo:
-        compiled_aten_module: torch.fx.GraphModule = torch_tensorrt.dynamo.compile(
+        compiled_aten_module: torch.fx.GraphModule = dynamo_compile(
             module,
-            inputs=inputs,
-            enabled_precisions=enabled_precisions,
+            inputs=input_list,
+            enabled_precisions=enabled_precisions_set,
             **kwargs,
         )
         return compiled_aten_module
     elif target_ir == _IRType.torch_compile:
-        return torch_compile(module, enabled_precisions=enabled_precisions, **kwargs)
+        return torch_compile(
+            module, enabled_precisions=enabled_precisions_set, **kwargs
+        )
     else:
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 
@@ -206,10 +219,10 @@ 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,
     ir: str = "default",
-    inputs: List[Input | torch.Tensor] = [],
-    enabled_precisions: Set[torch.dtype | dtype] = set([torch.float]),
+    enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,
 ) -> bytes:
     """Convert a TorchScript module method to a serialized TensorRT engine
@@ -242,6 +255,10 @@ def convert_method_to_trt_engine(
     Returns:
         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
     """
+    enabled_precisions_set = (
+        enabled_precisions if enabled_precisions is not None else {torch.float}
+    )
+
     module_type = _parse_module_type(module)
     target_ir = _get_target_ir(module_type, ir)
     if target_ir == _IRType.ts:
@@ -254,9 +271,9 @@ def convert_method_to_trt_engine(
             ts_mod = torch.jit.script(module)
         return torch_tensorrt.ts.convert_method_to_trt_engine(
             ts_mod,
-            method_name,
             inputs=inputs,
-            enabled_precisions=enabled_precisions,
+            method_name=method_name,
+            enabled_precisions=enabled_precisions_set,
             **kwargs,
         )
     elif target_ir == _IRType.fx:

py/torch_tensorrt/_enums.py

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

py/torch_tensorrt/dynamo/__init__.py

@@ -2,6 +2,4 @@
 from torch_tensorrt._util import sanitized_torch_version
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
-    from ._settings import *
-    from .compile import compile
-    from .aten_tracer import trace
+    from ._settings import *  # noqa: F403

py/torch_tensorrt/dynamo/_settings.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Sequence, Set
+from typing import Optional, Set
 import torch
 from torch_tensorrt.dynamo._defaults import (
     PRECISION,

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -2,7 +2,6 @@
 import sys
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
-from packaging import version
 
 import torch
 import torch._dynamo as torchdynamo

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -1 +0,0 @@
-from .backends import torch_tensorrt_backend

py/torch_tensorrt/dynamo/backend/backends.py

@@ -80,7 +80,7 @@ def _pretraced_backend(
             settings=settings,
         )
         return trt_compiled
-    except:
+    except AssertionError:
         if not settings.pass_through_build_failures:
             logger.warning(
                 "TRT conversion failed on the subgraph. See trace above. "

py/torch_tensorrt/dynamo/compile.py

@@ -2,12 +2,10 @@
 import logging
 import collections.abc
 import torch_tensorrt
-from functools import partial
 
-from typing import Any, Optional, Sequence, Set, List
+from typing import Any, Optional, Set, List, Tuple
 from torch_tensorrt import EngineCapability, Device
 from torch.fx.passes.pass_manager import PassManager
-from torch.fx.passes.shape_prop import ShapeProp
 from torch.fx.passes.splitter_base import SplitResult
 from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
 from torch_tensorrt.dynamo.lowering import (
@@ -16,7 +14,6 @@
 )
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
-from torch_tensorrt.dynamo.backend import torch_tensorrt_backend
 from torch_tensorrt.dynamo.backend.backends import _compile_module
 from torch_tensorrt.dynamo.conversion import convert_module
 
@@ -43,7 +40,7 @@ def compile(
     device: Device = Device._current_device(),
     disable_tf32: bool = False,
     sparse_weights: bool = False,
-    enabled_precisions: Set[torch.dtype] = set(),
+    enabled_precisions: Set[torch.dtype] | Tuple[torch.dtype] = (torch.float32,),
     refit: bool = False,
     debug: bool = DEBUG,
     capability: EngineCapability = EngineCapability.default,
@@ -56,8 +53,8 @@ def compile(
     truncate_long_and_double: bool = False,
     require_full_compilation: bool = False,
     min_block_size: int = MIN_BLOCK_SIZE,
-    torch_executed_ops: List[str] = [],
-    torch_executed_modules: List[str] = [],
+    torch_executed_ops: Optional[List[str]] = None,
+    torch_executed_modules: Optional[List[str]] = None,
     pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES,
     max_aux_streams: Optional[int] = MAX_AUX_STREAMS,
     version_compatible: bool = VERSION_COMPATIBLE,
@@ -68,7 +65,9 @@ def compile(
     if debug:
         logger.setLevel(logging.DEBUG)
 
-    logger.warn(
+    enabled_precisions = set(enabled_precisions)
+
+    logger.warning(
         "The Dynamo backend is an experimental feature, for which only the "
         + "following arguments are supported: "
         + "{enabled_precisions, debug, workspace_size, min_block_size, "
@@ -103,7 +102,9 @@ def compile(
         "debug": debug,
         "workspace_size": workspace_size,
         "min_block_size": min_block_size,
-        "torch_executed_ops": torch_executed_ops,
+        "torch_executed_ops": torch_executed_ops
+        if torch_executed_ops is not None
+        else [],
         "pass_through_build_failures": pass_through_build_failures,
         "max_aux_streams": max_aux_streams,
         "version_compatible": version_compatible,