fix: some more index fixes, enable dynamo tracer in some tests

Author: narendasan

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -6,7 +6,9 @@
 
 import numpy as np
 import torch
+from tensorrt import ITensor as TRTTensor
 from torch.fx.node import Argument, Node, Target
+
 from torch_tensorrt import ENABLED_FEATURES
 from torch_tensorrt._features import needs_not_tensorrt_rtx
 from torch_tensorrt._utils import is_tensorrt_version_supported
@@ -27,8 +29,6 @@
 )
 from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
 
-from tensorrt import ITensor as TRTTensor
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -2904,7 +2904,6 @@ def sort_validator(node: Node, settings: Optional[CompilationSettings] = None) -
 
 
 def topk_sort_validator(k: int) -> bool:
-
     # topk layer supports dynamic k value but we cannot determine supported dynamic topk value at
     # compile time.
     if k == DYNAMIC_DIM or not isinstance(k, int):

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -2,8 +2,11 @@
 from typing import List, Optional, Sequence, Union
 
 import numpy as np
+import tensorrt as trt
 import torch
+from tensorrt import ITensor
 from torch.fx.node import Target
+
 from torch_tensorrt._enums import dtype
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion import impl
@@ -21,9 +24,6 @@
 from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
 from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
 
-import tensorrt as trt
-from tensorrt import ITensor
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -875,9 +875,7 @@ def index_put_converter(
     K = len(I)
     # Determine the maximum size 'N' among the index tensors
     if K > 0:
-        index_shapes = (
-            []
-        )  # [tensor.shape[0] for tensor in indices if tensor is not None]
+        index_shapes = []  # [tensor.shape[0] for tensor in indices if tensor is not None]
         for _ni, idx_tensor in enumerate(indices):
             if idx_tensor is not None:
                 if idx_tensor.shape[0] != DYNAMIC_DIM:

tests/py/dynamo/conversion/test_index_aten.py

@@ -160,7 +160,7 @@ def forward(self, x, index0):
 
         input = torch.randn(2, 2)
         index0 = torch.tensor([True, False])
-        self.run_test(TestModule(), [input, index0], enable_passes=True)
+        self.run_test(TestModule(), [input, index0], use_dynamo_tracer=True, enable_passes=True)
 
     def test_index_zero_index_three_dim_ITensor(self):
         class TestModule(nn.Module):
@@ -172,7 +172,7 @@ def forward(self, x, index0):
         input = torch.randn(2, 2, 2)
         index0 = torch.randint(0, 1, (1, 1))
         index0 = index0.to(torch.int32)
-        self.run_test(TestModule(), [input, index0])
+        self.run_test(TestModule(), [input, index0], use_dynamo_tracer=True, enable_passes=True)
 
     @unittest.skipIf(
         ENABLED_FEATURES.tensorrt_rtx,
@@ -187,7 +187,7 @@ def forward(self, x, index0):
 
         input = torch.randn(2, 2, 2)
         index0 = torch.tensor([True, False])
-        self.run_test(TestModule(), [input, index0])
+        self.run_test(TestModule(), [input, index0], use_dynamo_tracer=True, enable_passes=True)
 
 
 class TestIndexDynamicConstantConverter(DispatchTestCase):