fix loops

Author: xadupre

_unittests/ut_export/test_control_flow.py renamed to _unittests/ut_export/test_control_flow_onnx.py

@@ -12,7 +12,7 @@
 from onnx_diagnostic.export.api import to_onnx
 
 
-class TestControlFlow(ExtTestCase):
+class TestControlFlowOnnx(ExtTestCase):
     @never_test()
     def test_loop_one_research(self):
         class Model(torch.nn.Module):
@@ -70,7 +70,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_one_custom_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlowOnnx_test_loop_one_custom_L_Model_forward_L_body_",
             str(ep),
         )
 
@@ -104,7 +104,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_one_custom_different_opset_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlowOnnx_test_loop_one_custom_different_opset_L_Model_forward_L_body_",
             str(ep),
         )
 
@@ -142,7 +142,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_two_custom_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlowOnnx_test_loop_two_custom_L_Model_forward_L_body_",
             str(ep),
         )
 
@@ -177,7 +177,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_two_custom_reduction_dim_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlowOnnx_test_loop_two_custom_reduction_dim_L_Model_forward_L_body_",
             str(ep),
         )
 

onnx_diagnostic/export/control_flow_onnx.py

@@ -265,7 +265,17 @@ def convert_custom_loop_into_onnx(
         nodes, graph.name, inputs, graph_outputs, graph.initializer, graph.sparse_initializer
     )
 
-    sequences = [g.op.SequenceEmpty() for _ in outputs]
+    itypes = [
+        graph.output[i].type.sequence_type.elem_type.tensor_type.elem_type
+        for i in range(1, len(graph.output))
+    ]
+    assert len(outputs) == len(
+        itypes
+    ), f"Length mismatch between outputs={outputs} and graph.output={graph.output}"
+    assert (
+        0 not in itypes
+    ), f"Undefined types are not allowed in itype={itypes}, graph.output={graph.output}"
+    sequences = [g.op.SequenceEmpty(dtype=itype) for itype in itypes]
 
     outloop = [g.unique_name(f"loop_for_onnx{i}") for i in range(len(sequences))]
 
@@ -285,8 +295,10 @@ def convert_custom_loop_into_onnx(
     ]
     if not sts:
         for i, o in enumerate(outputs):
-            g.set_type(o, graph.output[i].type.tensor_type.elem_type)
-            g.set_rank(o, len(graph.output[i].type.tensor_type.shape.dims))
+            g.set_type(o, graph.output[i].type.sequence_type.elem_type.tensor_type.elem_type)
+            g.set_rank(
+                o, len(graph.output[i].type.sequence_type.elem_type.tensor_type.shape.dims)
+            )
     return outputs if len(outputs) > 1 else outputs[0]
 
 

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -159,7 +159,9 @@ def append_output_sequence(
         """
         if not tensors:
             # empty list
-            self.nodes.append(oh.make_node("SequenceEmpty", [], [name]))
+            self.nodes.append(
+                oh.make_node("SequenceEmpty", [], [name], dtype=TensorProto.FLOAT)
+            )
             tensor_type_proto = oh.make_tensor_type_proto(
                 elem_type=TensorProto.FLOAT, shape=None
             )