feat: Add testing functionality to Python API utils

- Add function to check the equivalence of two collection-based outputs
for comparison across Torch-TRT and Torch outputs
- Improved test robustness in end-to-end to check for equivalent output
schemas in addition to successful compilation

Author: gs-olive

tests/py/api/test_e2e_behavior.py

@@ -4,6 +4,7 @@
 import torchvision.models as models
 import copy
 from typing import Dict
+from utils import same_output_format
 
 
 class TestInputTypeDefaultsFP32Model(unittest.TestCase):
@@ -109,7 +110,7 @@ def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
         )
         trt_mod(self.input)
 
-    def test_nested_tuple_output_with_full_compilation(self):
+    def test_nested_combination_tuple_list_output_with_full_compilation(self):
         class Sample(torch.nn.Module):
             def __init__(self):
                 super(Sample, self).__init__()
@@ -119,7 +120,7 @@ def forward(self, x, y, z):
                 b = x + 2.0 * z
                 b = y + b
                 a = b + c
-                return (a, (b, c))
+                return (a, [b, c])
 
         self.model = Sample().eval().to("cuda")
         self.input_1 = torch.zeros((5, 5), dtype=torch.float, device="cuda:0")
@@ -139,7 +140,11 @@ def forward(self, x, y, z):
             require_full_compilation=True,
             enabled_precisions={torch.float, torch.half},
         )
-        trt_mod(self.input_1, self.input_2, self.input_3)
+        trt_output = trt_mod(self.input_1, self.input_2, self.input_3)
+        torch_output = self.model(self.input_1, self.input_2, self.input_3)
+        assert same_output_format(
+            trt_output, torch_output
+        ), "Found differing output formatting between Torch-TRT and Torch"
 
 
 if __name__ == "__main__":

tests/py/api/utils.py

@@ -13,3 +13,42 @@ def cosine_similarity(gt_tensor, pred_tensor):
     res = res.cpu().detach().item()
 
     return res
+
+
+def same_output_format(trt_output, torch_output):
+    # For each encountered collection type, ensure the torch and trt outputs agree
+    # on type and size, checking recursively through all member elements.
+    if isinstance(trt_output, tuple):
+        return (
+            isinstance(torch_output, tuple)
+            and (len(trt_output) == len(torch_output))
+            and all(
+                same_output_format(trt_entry, torch_entry)
+                for trt_entry, torch_entry in zip(trt_output, torch_output)
+            )
+        )
+    elif isinstance(trt_output, list):
+        return (
+            isinstance(torch_output, list)
+            and (len(trt_output) == len(torch_output))
+            and all(
+                same_output_format(trt_entry, torch_entry)
+                for trt_entry, torch_entry in zip(trt_output, torch_output)
+            )
+        )
+    elif isinstance(trt_output, dict):
+        return (
+            isinstance(torch_output, dict)
+            and (len(trt_output) == len(torch_output))
+            and (trt_output.keys() == torch_output.keys())
+            and all(
+                same_output_format(trt_output[key], torch_output[key])
+                for key in trt_output.keys()
+            )
+        )
+    elif isinstance(trt_output, set) or isinstance(trt_output, frozenset):
+        raise AssertionError(
+            "Unsupported output type 'set' encountered in output format check."
+        )
+    else:
+        return type(trt_output) is type(torch_output)