Dynamo test cases error correction and adding support for int64 indices conversion to int32 in TRT10 for tests

Author: apbose

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

@@ -407,6 +407,12 @@ def scatter(
     index_shape = index.shape
     index_shape_list = list(index.shape)
     if not (isinstance(index, TRTTensor)):
+        if isinstance(index, torch.Tensor):
+            if index.dtype == torch.int64:
+                index = index.to(torch.int32)
+        elif isinstance(index, np.ndarray):
+            if index.dtype == np.int64:
+                index = index.astype(np.int32)
         index = get_trt_tensor(ctx, index, f"_index_tensor")
     dim = get_positive_dim(dim, len(input_shape))
     dynamic_shape = has_dynamic_shape(input.shape)

tests/py/dynamo/conversion/harness.py

@@ -226,6 +226,7 @@ def run_test(
         check_dtype=True,
         use_dynamo_tracer=False,
         enable_passes=False,
+        int32_reqd=False,
     ):
         mod.eval()
         mod = self.generate_graph(
@@ -245,10 +246,19 @@ def run_test(
 
         num_inputs = len(inputs)
         trt_inputs = inputs
+        dtype_to_change = []
+        if int32_reqd:
+            dtype_to_change = [torch.int64, torch.float64]
+        else:
+            dtype_to_change = [
+                torch.float64,
+            ]
         for num_input in range(num_inputs):
             input = inputs[num_input]
-            if input.dtype is torch.float64:
-                dtype_32bit = torch.float32
+            if input.dtype in dtype_to_change:
+                dtype_32bit = (
+                    torch.float32 if (input.dtype == torch.float64) else torch.int32
+                )
                 # should we modify graph here to insert clone nodes?
                 # ideally not required
                 trt_inputs = (
@@ -360,4 +370,4 @@ def run_test_with_dynamic_shape(
         # Since the lowering is based on optimal shape. We need to test with
         # different shape(for ex. max shape) for testing dynamic shape
         inputs_max = [spec.example_tensor("max_shape") for spec in input_specs]
-        super().run_test(mod, inputs_max, interp, rtol, atol)
+        super().run_test(mod, inputs_max, inputs_max, interp, rtol, atol)

tests/py/dynamo/conversion/test_scatter_aten.py

@@ -45,10 +45,7 @@ def forward(self, input):
 
         input = torch.zeros(3, 5, dtype=torch.int32)
         inputs = [input]
-        self.run_test(
-            TestModule(),
-            inputs,
-        )
+        self.run_test(TestModule(), inputs, int32_reqd=True)
 
     @parameterized.expand(
         [
@@ -78,10 +75,7 @@ def forward(self, input, index):
 
         input = torch.zeros(3, 5, dtype=torch.int32)
         inputs = [input, index]
-        self.run_test(
-            TestModule(),
-            inputs,
-        )
+        self.run_test(TestModule(), inputs, int32_reqd=True)
 
 
 class TestScatterSrcConverter(DispatchTestCase):
@@ -113,8 +107,18 @@ class TestScatterSrcConverter(DispatchTestCase):
             ),
             # These are special cases where in the harness.py code might need to be changed to input cuda_inputs
             # In that case below two test cases would also require index and src to be on cuda
-            # ("scatter_one_dim_indexOne_constant_src", 1, torch.tensor([[0, 1, 2, 0]]), torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)),
-            # ("scatter_one_dim_indexTwo_constant_src", 1, torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]), torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.int32)),
+            (
+                "scatter_one_dim_indexOne_constant_src",
+                1,
+                torch.tensor([[0, 1, 2, 0]]),
+                torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            ),
+            (
+                "scatter_one_dim_indexTwo_constant_src",
+                1,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+                torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.int32),
+            ),
         ]
     )
     def test_scatter_index_constant(self, _, dim, index, src):
@@ -128,10 +132,7 @@ def forward(self, input):
         input = torch.zeros(3, 5, dtype=torch.int32)
         inputs = [input]
         scatter = TestModule()
-        self.run_test(
-            TestModule(),
-            inputs,
-        )
+        self.run_test(TestModule(), inputs, int32_reqd=True)
 
     @parameterized.expand(
         [
@@ -171,10 +172,7 @@ def forward(self, input, index):
 
         input = torch.zeros(3, 5, dtype=torch.int32)
         inputs = [input, index]
-        self.run_test(
-            TestModule(),
-            inputs,
-        )
+        self.run_test(TestModule(), inputs, int32_reqd=True)
 
 
 if __name__ == "__main__":