Changing test cases to not get get_item trace and including test case for non continuous indices

Author: apbose

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

@@ -6,12 +6,14 @@
 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 broadcastable
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    broadcastable,
+    get_trt_tensor,
+)
 from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise
 from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.fx.converters.converter_utils import (
     get_positive_dim,
-    get_trt_tensor,
     has_dynamic_shape,
     set_layer_name,
     to_numpy,
@@ -88,6 +90,7 @@ def index(
 
     # here we need to check if all the index are broadcastable
     # if no, then we need to broadcast
+    input = get_trt_tensor(network, input, name + f"_input_to_tensor")
 
     last_index = None
     for i, ind in enumerate(index):
@@ -334,7 +337,7 @@ def index(
             concat_final_tensor = concat_final_shape_layer.get_output(0)
 
             reshape_layer = network.add_shuffle(gather_out)
-            reshape_layer.setInput(1, concat_final_tensor)
+            reshape_layer.set_input(1, concat_final_tensor)
             set_layer_name(
                 reshape_layer,
                 target,

tests/py/dynamo/conversion/test_index_aten.py

@@ -2,17 +2,22 @@
 
 import torch
 import torch.nn as nn
-from harness import DispatchTestCase
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
+from .harness import DispatchTestCase
+
 
 class TestIndexConverter(DispatchTestCase):
-    def test_index_zero(self):
+    def test_index_zero_two_dim(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.randint(0, 1, (1, 1))
+                super().__init__()
+
             def forward(self, x):
                 index0 = torch.randint(0, 1, (1, 1))
-                indices = [None, index0]
+                indices = [None, self.index0]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
@@ -23,11 +28,14 @@ def forward(self, x):
             expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one(self):
+    def test_index_zero_index_three_dim(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.randint(0, 1, (1, 1))
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.randint(0, 1, (1, 1))
-                indices = [None, index0, None]
+                indices = [None, self.index0, None]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
@@ -38,76 +46,101 @@ def forward(self, x):
             expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_index_two(self):
+    def test_index_zero_index_one_index_two_three_dim(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.randint(0, 1, (1, 1))
+                self.index1 = torch.randint(0, 1, (1, 1))
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.randint(0, 1, (1, 1))
-                index1 = torch.randint(0, 1, (1, 1))
-                indices = [None, index0, index1]
+                indices = [None, self.index0, self.index1]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
         input = [torch.randn(2, 2, 2)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
+            expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_SD(self):
+    def test_index_zero_index_one_four_dim(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.tensor([0, 0, 1, 1])
+                self.index1 = torch.tensor([0, 0, 1, 1])
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.tensor([0, 0, 1, 1])
-                index1 = torch.tensor([0, 0, 1, 1])
-                indices = [None, index0, index1, None]
+                indices = [None, self.index0, self.index1, None]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
         input = [torch.randn(2, 4, 4, 2)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
+            expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_SD(self):
+    def test_index_zero_index_one_four_dim_SD(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.tensor(
+                    [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]
+                )
+                self.index1 = torch.tensor(
+                    [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]
+                )
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                index1 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                indices = [None, index0, index1, None]
+                indices = [None, self.index0, self.index1, None]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
         input = [torch.randn(2, 1280, 8, 8)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
+            expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_SD(self):
+    def test_index_one_SD_unsqueeze_four_dim(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.tensor(
+                    [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]
+                )
+                self.index1 = self.index0.unsqueeze(0).T.long()
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                index1 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                indices = [None, index0, index1, None]
+                indices = [None, None, self.index1, self.index1]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
         input = [torch.randn(2, 1280, 8, 8)]
         self.run_test(
             TestModule(),
             input,
-            expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
+            expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_SD_unsqueeze(self):
+    def test_index_zero_index_one_index_two_SD_unsqueeze_four_dim_broadcast(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.tensor(
+                    [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]
+                )
+                self.index1 = self.index0.unsqueeze(0).T.long()
+                super().__init__()
+
             def forward(self, x):
                 index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
                 index1 = index0.unsqueeze(0).T.long()
-                indices = [None, None, index1, index1]
+                indices = [None, None, self.index0, self.index1]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
@@ -118,16 +151,19 @@ def forward(self, x):
             expected_ops={torch.ops.aten.index.Tensor},
         )
 
-    def test_index_zero_index_one_index_two_SD_unsqueeze(self):
+    def test_index_zero_index_one_index_four_dim_non_continuous(self):
         class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.tensor([0, 0, 1, 1])
+                self.index1 = torch.tensor([0, 0, 1, 1])
+                super().__init__()
+
             def forward(self, x):
-                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                index1 = index0.unsqueeze(0).T.long()
-                indices = [None, None, index0, index1]
+                indices = [None, self.index0, None, self.index1]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
 
-        input = [torch.randn(2, 1280, 8, 8)]
+        input = [torch.randn(2, 4, 4, 2)]
         self.run_test(
             TestModule(),
             input,