changing the test to compare size instead of elements

Author: apbose

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -73,7 +73,7 @@ def aten_ops_rand(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return np.random.rand(*args)
+    return np.random.rand(*args[0])
 
 
 @dynamo_tensorrt_converter(
@@ -86,7 +86,7 @@ def aten_ops_randn(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return np.random.randn(*args)
+    return np.random.randn(*args[0])
 
 
 def randperm_validator(randperm_node: Node) -> bool:
@@ -117,4 +117,4 @@ def aten_ops_randperm(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return np.random.permutation(*args)
+    return np.random.permutation(args[0])

tests/py/dynamo/conversion/harness.py

@@ -146,7 +146,6 @@ def run_test_custom_compare_results(
             res_trt = trt_mod(*cuda_inputs).cpu()
             res_cpu = mod(*inputs)
             assert len(res_trt) == len(res_cpu)
-            assert len(res_cpu) == len(comparators)
             for output_trt, output_cpu, comparator in zip(
                 res_trt, res_cpu, comparators
             ):
@@ -252,6 +251,40 @@ def run_test(
             check_dtype,
         )
 
+    def run_test_comparator(
+        self,
+        mod,
+        inputs,
+        expected_ops,
+        comparators: List[Tuple[Callable, List]],
+        precision=torch.float32,
+        output_dtypes=None,
+        use_dynamo_tracer=False,
+        enable_passes=False,
+    ):
+        mod.eval()
+        mod = self.generate_graph(
+            mod,
+            inputs,
+            use_dynamo_tracer=use_dynamo_tracer,
+            enable_passes=enable_passes,
+        )
+        # Previous instance of the interpreter auto-casted 64-bit inputs
+        # We replicate this behavior here
+        compilation_settings = CompilationSettings(
+            precision=precision, truncate_long_and_double=True
+        )
+
+        interp = TRTInterpreter(
+            mod,
+            Input.from_tensors(inputs),
+            output_dtypes=output_dtypes,
+            compilation_settings=compilation_settings,
+        )
+        super().run_test_custom_compare_results(
+            mod, inputs, expected_ops, interp, comparators
+        )
+
     def run_test_with_dynamic_shape(
         self,
         mod,

tests/py/dynamo/conversion/test_rand_aten.py

@@ -4,6 +4,8 @@
 from parameterized import parameterized
 from torch.testing._internal.common_utils import TestCase, run_tests
 
+from .harness import DispatchTestCase
+
 rand_ops = [
     (
         "rand_one_dimension",
@@ -13,7 +15,7 @@
     (
         "rand_two_dimension",
         (lambda shape: torch.ops.aten.rand(shape)),
-        [2, 3],
+        [1, 2],
     ),
     (
         "rand_three_dimension",
@@ -35,25 +37,29 @@
         (lambda shape: torch.ops.aten.randn(shape)),
         [2, 3, 4],
     ),
+]
+
+
+rand_perm_ops = [
     (
         "randperm_one_case",
         (lambda x: torch.ops.aten.randperm(x)),
-        1,
+        [1],
     ),
     (
         "randperm_two_case",
         (lambda x: torch.ops.aten.randperm(x)),
-        150,
+        [150],
     ),
     (
         "randperm_three_case",
         (lambda x: torch.ops.aten.randperm(x)),
-        1500,
+        [1500],
     ),
 ]
 
 
-class TestRandConverter(TestCase):
+class TestRandConverter(DispatchTestCase):
     @parameterized.expand(
         [
             (
@@ -64,41 +70,64 @@ class TestRandConverter(TestCase):
             for rand_op in rand_ops
         ]
     )
-    def test_rand(self, _, op, shape_or_input):
+    def test_rand(self, name, op, shape_or_input):
         class TestModule(nn.Module):
-            def __init__(self, rand_op, size):
+            def __init__(self):
                 super().__init__()
-                self.rand_op = rand_op
-                self.size = size
 
-            def forward(self):
-                return self.rand_op(self.size)
+            def forward(self, x):
+                shape_or_input[0] = x.shape[0]
+                return op(shape_or_input)
 
-        rand_model = TestModule(op, shape_or_input)
-        # cannot use self.run_test() since it expects input in form of tensor
+        rand_model = TestModule()
 
-        fx_graph = torch.fx.symbolic_trace(grid_model)
-        torch._dynamo.reset()
-
-        optimized_model = torch_tensorrt.compile(
-            fx_graph,
-            "torch_compile",
-            None,
-            min_block_size=1,
-            pass_through_build_failures=True,
-            truncate_long_and_double=True,
-            debug=True,
+        inputs = [torch.randint(1, 3, shape_or_input, dtype=torch.int32)]
+        comparator_shape = lambda x, y, check_dtype: x.shape == y.shape and (
+            x.dtype == y.dtype if check_dtype else True
+        )
+        expected_ops = []
+        self.run_test_comparator(
+            rand_model,
+            inputs,
+            expected_ops,
+            [(comparator_shape, [True])],
+            use_dynamo_tracer=True,
         )
-        optimized_model_results = optimized_model().detach().cpu()
-        torch_model_results = fx_graph().detach().cpu()
-        max_diff = float(
-            torch.max(torch.abs(optimized_model_results - torch_model_results))
+
+    @parameterized.expand(
+        [
+            (
+                rand_op[0],
+                rand_op[1],
+                rand_op[2],
+            )
+            for rand_op in rand_perm_ops
+        ]
+    )
+    def test_rand(self, name, op, shape_or_input):
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                shape_or_input[0] = x.shape[0]
+                return op(shape_or_input[0])
+
+        rand_model = TestModule()
+        # cannot use self.run_test() since it expects input in form of tensor
+
+        inputs = [torch.randint(1, 3, shape_or_input, dtype=torch.int32)]
+        comparator_shape = lambda x, y, check_dtype: x.shape == y.shape and (
+            x.dtype == y.dtype if check_dtype else True
         )
-        self.assertAlmostEqual(
-            max_diff,
-            0,
-            4,
-            f"TRT outputs don't match with the original model.",
+        expected_ops = []
+        # TRT-np returns int32 while torch returns float32
+        self.run_test_comparator(
+            rand_model,
+            inputs,
+            expected_ops,
+            [(comparator_shape, [False])],
+            use_dynamo_tracer=True,
         )