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

core/compiler.cpp

@@ -196,9 +196,8 @@ partitioning::GraphAndMapping BuildHybridGraph(
         // for collections processing
         if (expect_full_compilation) {
           for (auto torch_node : seg_block.block()->nodes()) {
-            if (partitioning::CollectionSchemas.find(torch_node->kind().toQualString()) ==
-                partitioning::CollectionSchemas.end()) {
-              LOG_WARNING(
+            if (partitioning::CollectionNodeKinds.find(torch_node->kind()) == partitioning::CollectionNodeKinds.end()) {
+              LOG_ERROR(
                   "Full compilation specified but node " << torch_node->kind().toQualString()
                                                          << " was executed in Torch.");
             }
@@ -210,7 +209,7 @@ partitioning::GraphAndMapping BuildHybridGraph(
     // If full compilation is expected, cannot have more than 2 Torch segments
     // (one for preprocessing inputs, one for post-processing outputs) and 1 TRT segment
     if (expect_full_compilation && !(num_torch_segments <= 2 && num_trt_segments == 1)) {
-      LOG_WARNING(
+      LOG_ERROR(
           "Full compilation specified but number of torch segments was "
           << num_torch_segments << " and number of trt segments was " << num_trt_segments
           << ". Was expecting at most 2 Torch segments and 1 TRT segment.");

core/partitioning/partitioning.h

@@ -21,14 +21,14 @@ typedef std::pair<std::shared_ptr<torch::jit::Graph>, std::unordered_map<torch::
 // Set of schemas allowed to be executed in Torch, even with require_full_compilation=true,
 // as necessary for returning collections of Tensors or other complex constructs, and for
 // processing inputs to TRT engines
-const std::unordered_set<std::string> CollectionSchemas = {
-    "prim::Constant",
-    "aten::__getitem__",
-    "prim::ListConstruct",
-    "prim::ListUnpack",
-    "prim::TupleIndex",
-    "prim::TupleConstruct",
-    "prim::TupleUnpack",
+const std::unordered_set<c10::Symbol> CollectionNodeKinds = {
+    c10::Symbol::fromQualString("prim::Constant"),
+    c10::Symbol::fromQualString("aten::__getitem__"),
+    c10::Symbol::fromQualString("prim::ListConstruct"),
+    c10::Symbol::fromQualString("prim::ListUnpack"),
+    c10::Symbol::fromQualString("prim::TupleIndex"),
+    c10::Symbol::fromQualString("prim::TupleConstruct"),
+    c10::Symbol::fromQualString("prim::TupleUnpack"),
 };
 
 ExampleIValues generateRandomInputs(

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)