From e1c9c43017d723ca0fd0b701083e237e14871349 Mon Sep 17 00:00:00 2001 From: HolyWu Date: Wed, 29 Jan 2025 19:48:55 +0800 Subject: [PATCH 1/2] Use IUnsqueezeLayer in unsqueeze impl --- .../dynamo/conversion/aten_ops_converters.py | 10 +-- .../dynamo/conversion/impl/unsqueeze.py | 68 ++----------------- 2 files changed, 13 insertions(+), 65 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index eac8a4b70c..4f2d168d29 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -7,7 +7,6 @@ import numpy as np import torch from torch.fx.node import Argument, Node, Target - from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl @@ -650,6 +649,11 @@ def aten_ops_erf( @dynamo_tensorrt_converter( torch.ops.aten.unsqueeze.default, supports_dynamic_shapes=True ) +@enforce_tensor_types( + { + 0: (TRTTensor,), + } +) def aten_ops_unsqueeze( ctx: ConversionContext, target: Target, @@ -657,9 +661,7 @@ def aten_ops_unsqueeze( kwargs: Dict[str, Argument], name: str, ) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.unsqueeze.unsqueeze( - ctx, target, SourceIR.ATEN, name, input_t=args[0], dim=args[1] - ) + return impl.unsqueeze.unsqueeze(ctx, target, SourceIR.ATEN, name, args[0], args[1]) @dynamo_tensorrt_converter( diff --git a/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py b/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py index 0db246a7c6..3dacc2fbe4 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py @@ -1,14 +1,13 @@ -from typing import List, Optional, Sequence, cast +from typing import List, Optional, Sequence from torch.fx.node import Target from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext from torch_tensorrt.dynamo.conversion.converter_utils import ( - get_positive_dim, get_trt_tensor, + set_layer_name, ) -from torch_tensorrt.fx.converters.converter_utils import set_layer_name -from torch_tensorrt.fx.types import Shape, TRTTensor +from torch_tensorrt.dynamo.types import TRTTensor def unsqueeze( @@ -16,64 +15,11 @@ def unsqueeze( target: Target, source_ir: Optional[SourceIR], name: str, - input_t: TRTTensor, - dim: Shape, + input: TRTTensor, + dim: int, ) -> TRTTensor: - input_val = get_trt_tensor(ctx, input_t, f"{name}_input_t") - if not isinstance(input_val, TRTTensor): - raise RuntimeError( - f"unsqueeze received input {input_val} that is not part " - "of the TensorRT region!" - ) - - dim = cast(int, dim) - - input_shape_size = len(input_val.shape) - dim = get_positive_dim(dim, input_shape_size + 1) - - intermediate_dim = 0 - dynamic_shape_cnt = 0 - # if unsqueeze the last dimensions, we can directly append to the shape - if dim == input_shape_size: - intermediate_dim = dim - else: - # since maximum of one dimension is permitted to be specified as -1 - # find the intermediate_dim which has only 1 dynamic_shape_cnt - # and then we can add a transpose after reshape if it is not the final shape we want - for i, s in reversed(list(enumerate(input_val.shape))): - if i >= dim: - if s == -1: - dynamic_shape_cnt += 1 - if dynamic_shape_cnt > 1: - intermediate_dim = i + 1 - break - if i == dim: - intermediate_dim = i - break - # calculate the new_shape for the shuffle layer's reshape_dims - new_shape = list( - tuple(input_val.shape)[:intermediate_dim] - + (1,) - + tuple(input_val.shape)[intermediate_dim:] - ) - for i, s in enumerate(new_shape): - if i < intermediate_dim and s == -1: - new_shape[i] = 0 - layer = ctx.net.add_shuffle(input_val) - layer.reshape_dims = tuple(new_shape) - # if the intermediate_dim is not the final dim we want to unsqueeze, add a second_transpose after reshape - if intermediate_dim != dim: - # calculate the second_transpose for the shuffle layer - permutation = [*range(0, len(new_shape))] - # for example: if the reshape_dims is (3, 3, 5, 1, 5) and the final shape we want is (3, 1, 3, 5, 5) - # here intermediate_dim=3, dim=1, we need to move intermediate_dim before [dim: intermediate_dim) - new_permutation = ( - tuple(permutation[:dim]) - + (intermediate_dim,) - + tuple(permutation[dim:intermediate_dim]) - + tuple(permutation[intermediate_dim + 1 :]) - ) - layer.second_transpose = new_permutation + axes = get_trt_tensor(ctx, dim, f"{name}_axes") + layer = ctx.net.add_unsqueeze(input, axes) set_layer_name(layer, target, name, source_ir) return layer.get_output(0) From f496605561f98a1bd49ff70bd1a59c3e6375822e Mon Sep 17 00:00:00 2001 From: HolyWu Date: Thu, 30 Jan 2025 10:42:42 +0800 Subject: [PATCH 2/2] black format --- .../torch_compile_advanced_usage.py | 3 ++- .../2a9ac10f2667047a7f398d1593b7ca33/torch_export_gpt2.py | 3 ++- .../torch_export_cudagraphs.py | 3 ++- .../dynamo_compile_resnet_example.py | 3 ++- .../7b7004dc2ea6f839be532665e16e0426/torch_export_llama2.py | 3 ++- .../torch_compile_resnet_example.py | 3 ++- .../torch_compile_transformers_example.py | 3 ++- .../dynamo_compile_advanced_usage.py | 3 ++- .../dynamo_compile_transformers_example.py | 3 ++- .../dynamo_compile_resnet_example.py | 3 ++- .../dynamo_compile_advanced_usage.py | 3 ++- .../dynamo_compile_transformers_example.py | 3 ++- examples/dynamo/torch_compile_advanced_usage.py | 3 ++- examples/dynamo/torch_compile_resnet_example.py | 3 ++- examples/dynamo/torch_compile_transformers_example.py | 3 ++- examples/dynamo/torch_export_cudagraphs.py | 3 ++- examples/dynamo/torch_export_gpt2.py | 3 ++- examples/dynamo/torch_export_llama2.py | 3 ++- py/torch_tensorrt/_Input.py | 2 +- py/torch_tensorrt/_enums.py | 2 +- py/torch_tensorrt/dynamo/conversion/_TRTBuilderMonitor.py | 6 +++--- py/torch_tensorrt/dynamo/conversion/impl/activation/ops.py | 4 ++-- py/torch_tensorrt/dynamo/utils.py | 2 +- py/torch_tensorrt/fx/test/converters/acc_op/test_where.py | 2 +- py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py | 2 +- 25 files changed, 46 insertions(+), 28 deletions(-) diff --git a/docs/_downloads/0e30a6276601af7e5fc4d5166e2e3d37/torch_compile_advanced_usage.py b/docs/_downloads/0e30a6276601af7e5fc4d5166e2e3d37/torch_compile_advanced_usage.py index 8ebedab111..af7d4b212d 100644 --- a/docs/_downloads/0e30a6276601af7e5fc4d5166e2e3d37/torch_compile_advanced_usage.py +++ b/docs/_downloads/0e30a6276601af7e5fc4d5166e2e3d37/torch_compile_advanced_usage.py @@ -4,7 +4,8 @@ Torch Compile Advanced Usage ====================================================== -This interactive script is intended as an overview of the process by which `torch_tensorrt.compile(..., ir="torch_compile", ...)` works, and how it integrates with the `torch.compile` API.""" +This interactive script is intended as an overview of the process by which `torch_tensorrt.compile(..., ir="torch_compile", ...)` works, and how it integrates with the `torch.compile` API. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/2a9ac10f2667047a7f398d1593b7ca33/torch_export_gpt2.py b/docs/_downloads/2a9ac10f2667047a7f398d1593b7ca33/torch_export_gpt2.py index cea0f3adf2..4d34c58de4 100644 --- a/docs/_downloads/2a9ac10f2667047a7f398d1593b7ca33/torch_export_gpt2.py +++ b/docs/_downloads/2a9ac10f2667047a7f398d1593b7ca33/torch_export_gpt2.py @@ -4,7 +4,8 @@ Compiling GPT2 using the dynamo backend ========================================================== -This script illustrates Torch-TensorRT workflow with dynamo backend on popular GPT2 model.""" +This script illustrates Torch-TensorRT workflow with dynamo backend on popular GPT2 model. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/3d4d74f6636d986f33167154f6553961/torch_export_cudagraphs.py b/docs/_downloads/3d4d74f6636d986f33167154f6553961/torch_export_cudagraphs.py index 1671c7783d..fb31766b7c 100644 --- a/docs/_downloads/3d4d74f6636d986f33167154f6553961/torch_export_cudagraphs.py +++ b/docs/_downloads/3d4d74f6636d986f33167154f6553961/torch_export_cudagraphs.py @@ -4,7 +4,8 @@ Torch Export with Cudagraphs ====================================================== -This interactive script is intended as an overview of the process by which the Torch-TensorRT Cudagraphs integration can be used in the `ir="dynamo"` path. The functionality works similarly in the `torch.compile` path as well.""" +This interactive script is intended as an overview of the process by which the Torch-TensorRT Cudagraphs integration can be used in the `ir="dynamo"` path. The functionality works similarly in the `torch.compile` path as well. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py b/docs/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py index 797e41f5fd..5826e28d1e 100644 --- a/docs/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py +++ b/docs/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py @@ -4,7 +4,8 @@ Compiling ResNet using the Torch-TensorRT Dyanmo Frontend ========================================================== -This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a ResNet model.""" +This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a ResNet model. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/7b7004dc2ea6f839be532665e16e0426/torch_export_llama2.py b/docs/_downloads/7b7004dc2ea6f839be532665e16e0426/torch_export_llama2.py index 5cfd1ed61c..2f3e3cba43 100644 --- a/docs/_downloads/7b7004dc2ea6f839be532665e16e0426/torch_export_llama2.py +++ b/docs/_downloads/7b7004dc2ea6f839be532665e16e0426/torch_export_llama2.py @@ -4,7 +4,8 @@ Compiling Llama2 using the dynamo backend ========================================================== -This script illustrates Torch-TensorRT workflow with dynamo backend on popular Llama2 model.""" +This script illustrates Torch-TensorRT workflow with dynamo backend on popular Llama2 model. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/d6e1bb6ec5f884994554d9d12e37a0f6/torch_compile_resnet_example.py b/docs/_downloads/d6e1bb6ec5f884994554d9d12e37a0f6/torch_compile_resnet_example.py index f852d60158..fb75986099 100644 --- a/docs/_downloads/d6e1bb6ec5f884994554d9d12e37a0f6/torch_compile_resnet_example.py +++ b/docs/_downloads/d6e1bb6ec5f884994554d9d12e37a0f6/torch_compile_resnet_example.py @@ -4,7 +4,8 @@ Compiling ResNet with dynamic shapes using the `torch.compile` backend ========================================================== -This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a ResNet model.""" +This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a ResNet model. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/dfa60e8f9850fd7761f3e7da81304d32/torch_compile_transformers_example.py b/docs/_downloads/dfa60e8f9850fd7761f3e7da81304d32/torch_compile_transformers_example.py index 221ecd4fd1..17cf46e8a3 100644 --- a/docs/_downloads/dfa60e8f9850fd7761f3e7da81304d32/torch_compile_transformers_example.py +++ b/docs/_downloads/dfa60e8f9850fd7761f3e7da81304d32/torch_compile_transformers_example.py @@ -4,7 +4,8 @@ Compiling BERT using the `torch.compile` backend ============================================================== -This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a BERT model.""" +This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a BERT model. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py b/docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py index f73bd1e780..3fb63e8a32 100644 --- a/docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py +++ b/docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py @@ -4,7 +4,8 @@ Dynamo Compile Advanced Usage ====================================================== -This interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API.""" +This interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API. +""" # %% # Imports and Model Definition diff --git a/docs/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py b/docs/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py index dd7fe2e07a..59319078a4 100644 --- a/docs/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py +++ b/docs/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py @@ -4,7 +4,8 @@ Compiling a Transformer using torch.compile and TensorRT ============================================================== -This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a transformer-based model.""" +This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a transformer-based model. +""" # %% # Imports and Model Definition diff --git a/docs/v1.4.0/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py b/docs/v1.4.0/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py index 797e41f5fd..5826e28d1e 100644 --- a/docs/v1.4.0/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py +++ b/docs/v1.4.0/_downloads/418941399c146271a7b7728ba3059960/dynamo_compile_resnet_example.py @@ -4,7 +4,8 @@ Compiling ResNet using the Torch-TensorRT Dyanmo Frontend ========================================================== -This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a ResNet model.""" +This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a ResNet model. +""" # %% # Imports and Model Definition diff --git a/docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py b/docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py index f73bd1e780..3fb63e8a32 100644 --- a/docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py +++ b/docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py @@ -4,7 +4,8 @@ Dynamo Compile Advanced Usage ====================================================== -This interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API.""" +This interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API. +""" # %% # Imports and Model Definition diff --git a/docs/v1.4.0/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py b/docs/v1.4.0/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py index dd7fe2e07a..59319078a4 100644 --- a/docs/v1.4.0/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py +++ b/docs/v1.4.0/_downloads/e550c5f53cc43e11aa6da8cfb79b54df/dynamo_compile_transformers_example.py @@ -4,7 +4,8 @@ Compiling a Transformer using torch.compile and TensorRT ============================================================== -This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a transformer-based model.""" +This interactive script is intended as a sample of the `torch_tensorrt.dynamo.compile` workflow on a transformer-based model. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_compile_advanced_usage.py b/examples/dynamo/torch_compile_advanced_usage.py index 8ebedab111..af7d4b212d 100644 --- a/examples/dynamo/torch_compile_advanced_usage.py +++ b/examples/dynamo/torch_compile_advanced_usage.py @@ -4,7 +4,8 @@ Torch Compile Advanced Usage ====================================================== -This interactive script is intended as an overview of the process by which `torch_tensorrt.compile(..., ir="torch_compile", ...)` works, and how it integrates with the `torch.compile` API.""" +This interactive script is intended as an overview of the process by which `torch_tensorrt.compile(..., ir="torch_compile", ...)` works, and how it integrates with the `torch.compile` API. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_compile_resnet_example.py b/examples/dynamo/torch_compile_resnet_example.py index f852d60158..fb75986099 100644 --- a/examples/dynamo/torch_compile_resnet_example.py +++ b/examples/dynamo/torch_compile_resnet_example.py @@ -4,7 +4,8 @@ Compiling ResNet with dynamic shapes using the `torch.compile` backend ========================================================== -This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a ResNet model.""" +This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a ResNet model. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_compile_transformers_example.py b/examples/dynamo/torch_compile_transformers_example.py index 221ecd4fd1..17cf46e8a3 100644 --- a/examples/dynamo/torch_compile_transformers_example.py +++ b/examples/dynamo/torch_compile_transformers_example.py @@ -4,7 +4,8 @@ Compiling BERT using the `torch.compile` backend ============================================================== -This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a BERT model.""" +This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a BERT model. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_export_cudagraphs.py b/examples/dynamo/torch_export_cudagraphs.py index 1671c7783d..fb31766b7c 100644 --- a/examples/dynamo/torch_export_cudagraphs.py +++ b/examples/dynamo/torch_export_cudagraphs.py @@ -4,7 +4,8 @@ Torch Export with Cudagraphs ====================================================== -This interactive script is intended as an overview of the process by which the Torch-TensorRT Cudagraphs integration can be used in the `ir="dynamo"` path. The functionality works similarly in the `torch.compile` path as well.""" +This interactive script is intended as an overview of the process by which the Torch-TensorRT Cudagraphs integration can be used in the `ir="dynamo"` path. The functionality works similarly in the `torch.compile` path as well. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_export_gpt2.py b/examples/dynamo/torch_export_gpt2.py index cea0f3adf2..4d34c58de4 100644 --- a/examples/dynamo/torch_export_gpt2.py +++ b/examples/dynamo/torch_export_gpt2.py @@ -4,7 +4,8 @@ Compiling GPT2 using the dynamo backend ========================================================== -This script illustrates Torch-TensorRT workflow with dynamo backend on popular GPT2 model.""" +This script illustrates Torch-TensorRT workflow with dynamo backend on popular GPT2 model. +""" # %% # Imports and Model Definition diff --git a/examples/dynamo/torch_export_llama2.py b/examples/dynamo/torch_export_llama2.py index 5cfd1ed61c..2f3e3cba43 100644 --- a/examples/dynamo/torch_export_llama2.py +++ b/examples/dynamo/torch_export_llama2.py @@ -4,7 +4,8 @@ Compiling Llama2 using the dynamo backend ========================================================== -This script illustrates Torch-TensorRT workflow with dynamo backend on popular Llama2 model.""" +This script illustrates Torch-TensorRT workflow with dynamo backend on popular Llama2 model. +""" # %% # Imports and Model Definition diff --git a/py/torch_tensorrt/_Input.py b/py/torch_tensorrt/_Input.py index 126219ee8a..2f953094ca 100644 --- a/py/torch_tensorrt/_Input.py +++ b/py/torch_tensorrt/_Input.py @@ -261,7 +261,7 @@ def _supported_input_size_type(input_size: Any) -> bool: @staticmethod def _parse_tensor_domain( - domain: Optional[Tuple[float, float]] + domain: Optional[Tuple[float, float]], ) -> Tuple[float, float]: """ Produce a tuple of integers which specifies a tensor domain in the interval format: [lo, hi) diff --git a/py/torch_tensorrt/_enums.py b/py/torch_tensorrt/_enums.py index eaefb68ce5..c706c345d6 100644 --- a/py/torch_tensorrt/_enums.py +++ b/py/torch_tensorrt/_enums.py @@ -1200,7 +1200,7 @@ def _from( @classmethod def try_from( - c: Union[trt.EngineCapability, EngineCapability] + c: Union[trt.EngineCapability, EngineCapability], ) -> Optional[EngineCapability]: """Create a Torch-TensorRT engine capability enum from a TensorRT engine capability enum. diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTBuilderMonitor.py b/py/torch_tensorrt/dynamo/conversion/_TRTBuilderMonitor.py index 9a1189e44a..9b2755f4c7 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTBuilderMonitor.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTBuilderMonitor.py @@ -53,13 +53,13 @@ def _redraw(self, *, blank_lines: int = 0) -> None: if self._render: def clear_line() -> None: - print("\x1B[2K", end="") + print("\x1b[2K", end="") def move_to_start_of_line() -> None: - print("\x1B[0G", end="") + print("\x1b[0G", end="") def move_cursor_up(lines: int) -> None: - print("\x1B[{}A".format(lines), end="") + print("\x1b[{}A".format(lines), end="") def progress_bar(steps: int, num_steps: int) -> str: INNER_WIDTH = 10 diff --git a/py/torch_tensorrt/dynamo/conversion/impl/activation/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/activation/ops.py index a563118526..eb981f2031 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/activation/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/activation/ops.py @@ -247,7 +247,7 @@ def hard_sigmoid( operation_type = trt.ActivationType.HARD_SIGMOID def hard_sigmoid_dyn_range_fn( - dyn_range: Tuple[float, float] + dyn_range: Tuple[float, float], ) -> Tuple[float, float]: def hard_sigmoid_fn(x: float) -> float: return max(0, min(1, alpha * x + beta)) @@ -310,7 +310,7 @@ def thresholded_relu( operation_type = trt.ActivationType.THRESHOLDED_RELU def thresholded_relu_dyn_range_fn( - dyn_range: Tuple[float, float] + dyn_range: Tuple[float, float], ) -> Tuple[float, float]: def thresholded_relu_fn(x: float) -> float: return x if x > alpha else 0 diff --git a/py/torch_tensorrt/dynamo/utils.py b/py/torch_tensorrt/dynamo/utils.py index 5d6807f33a..6e5bece418 100644 --- a/py/torch_tensorrt/dynamo/utils.py +++ b/py/torch_tensorrt/dynamo/utils.py @@ -465,7 +465,7 @@ def to_torch_device(device: Optional[Union[Device, torch.device, str]]) -> torch def to_torch_tensorrt_device( - device: Optional[Union[Device, torch.device, str]] + device: Optional[Union[Device, torch.device, str]], ) -> Device: """Cast a device-type to torch_tensorrt.Device diff --git a/py/torch_tensorrt/fx/test/converters/acc_op/test_where.py b/py/torch_tensorrt/fx/test/converters/acc_op/test_where.py index 72fea70265..1e14b50305 100644 --- a/py/torch_tensorrt/fx/test/converters/acc_op/test_where.py +++ b/py/torch_tensorrt/fx/test/converters/acc_op/test_where.py @@ -101,7 +101,7 @@ def __init__(self, x_shape, y_shape): def forward(self, condition): return torch.where(condition, self.x, self.y) - inputs = [(torch.randn(condition_shape) > 0)] + inputs = [torch.randn(condition_shape) > 0] self.run_test( Where(x_shape, y_shape), inputs, diff --git a/py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py b/py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py index 9d5576bd63..677fd96095 100644 --- a/py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py +++ b/py/torch_tensorrt/fx/tracer/acc_tracer/acc_tracer.py @@ -517,7 +517,7 @@ def _replace_transpose_last_dims_impl( changed = False def _calculate_dim( - transpose_dim: Union[torch.fx.Node, int] + transpose_dim: Union[torch.fx.Node, int], ) -> Union[torch.fx.Node, int]: nonlocal transpose_input_node nonlocal changed