Flatten all the output tensors to a list when the returned outputs contain list of list.

Author: Xreki

graph_net/paddle/test_compiler.py

@@ -221,47 +221,56 @@ def measure_performance(model_call, args, compiler, profile=False):
 
 
 def check_outputs(args, expected_out, compiled_out):
-    if isinstance(expected_out, paddle.Tensor):
-        expected_out = [expected_out]
-    if isinstance(compiled_out, paddle.Tensor):
-        compiled_out = [compiled_out]
-
-    eager_dtypes = [None] * len(expected_out)
-    for i, tensor in enumerate(expected_out):
-        eager_dtypes[i] = (
-            str(tensor.dtype).replace("paddle.", "") if tensor is not None else "None"
-        )
-
-    compiled_dtypes = [None] * len(compiled_out)
-    for i, tensor in enumerate(compiled_out):
-        compiled_dtypes[i] = (
-            str(tensor.dtype).replace("paddle.", "") if tensor is not None else "None"
-        )
-
+    def _flatten_outputs_to_list(outs):
+        flattened_outs = outs
+        if isinstance(outs, paddle.Tensor):
+            flattened_outs = [outs]
+        else:
+            flattened_outs = [
+                x
+                for out in outs
+                for x in (out if isinstance(out, (tuple, list)) else (out,))
+            ]
+        return flattened_outs
+
+    expected_out = _flatten_outputs_to_list(expected_out)
+    compiled_out = _flatten_outputs_to_list(compiled_out)
+
+    def _get_output_dtypes(outs):
+        dtypes = [
+            str(tensor.dtype).replace("paddle.", "")
+            if isinstance(tensor, paddle.Tensor)
+            else None
+            for i, tensor in enumerate(outs)
+        ]
+        return dtypes
+
+    eager_dtypes = _get_output_dtypes(expected_out)
+    compiled_dtypes = _get_output_dtypes(compiled_out)
     type_match = test_compiler_util.check_output_datatype(
         args, eager_dtypes, compiled_dtypes
     )
 
-    eager_shapes = [None] * len(expected_out)
-    for i, tensor in enumerate(expected_out):
-        eager_shapes[i] = tensor.shape if tensor is not None else None
-
-    compiled_shapes = [None] * len(compiled_out)
-    for i, tensor in enumerate(compiled_out):
-        compiled_shapes[i] = tensor.shape if tensor is not None else None
+    def _get_output_shapes(outs):
+        shapes = [
+            tensor.shape if isinstance(tensor, paddle.Tensor) else None
+            for i, tensor in enumerate(outs)
+        ]
+        return shapes
 
+    eager_shapes = _get_output_shapes(expected_out)
+    compiled_shapes = _get_output_shapes(compiled_out)
     shape_match = test_compiler_util.check_output_shape(
         args, eager_shapes, compiled_shapes
     )
 
-    def transfer_to_float(origin_outputs):
+    def _transfer_to_float(origin_outputs):
         outputs = []
         for item in origin_outputs:
-            if (
-                item is not None
-                and isinstance(item, paddle.Tensor)
-                and item.dtype not in [paddle.float32, paddle.float64]
-            ):
+            if isinstance(item, paddle.Tensor) and item.dtype not in [
+                paddle.float32,
+                paddle.float64,
+            ]:
                 item = item.astype("float32")
             outputs.append(item)
         return outputs
@@ -274,8 +283,8 @@ def transfer_to_float(origin_outputs):
             cmp_equal_func=get_cmp_equal,
         )
 
-        expected_out_fp32 = transfer_to_float(expected_out)
-        compiled_out_fp32 = transfer_to_float(compiled_out)
+        expected_out_fp32 = _transfer_to_float(expected_out)
+        compiled_out_fp32 = _transfer_to_float(compiled_out)
         test_compiler_util.check_allclose(
             args,
             expected_out_fp32,