Update parent graph in loop "cond" subgraphs (#2201)

* handle while cond graph subgraph parent update
* add while loop cond subgrpah test

---------

Signed-off-by: Salvetti, Francesco <[email protected]>

Author: f-salvetti

tests/test_loops.py

@@ -8,7 +8,7 @@
 
 from backend_test_base import Tf2OnnxBackendTestBase
 from common import unittest_main, check_tf_min_version, check_tf_max_version, \
-    check_onnxruntime_min_version, check_tfjs_max_version
+    check_onnxruntime_min_version, check_tfjs_max_version, skip_tflite
 from tf2onnx.tf_loader import is_tf2
 
 
@@ -302,6 +302,23 @@ def func(i):
         output_names_with_port = ["output:0"]
         self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)
 
+    @check_tf_min_version("2")
+    @skip_tflite("cond_graph conversion fails with tflite")
+    def test_while_loop_cond_subgraphs(self):
+        # test for while_loop with subgraphs in cond
+        # Note: this is not working on tf1
+        def func(x):
+            x_dim = tf.shape(x)[0]
+            r = tf.cast(tf.zeros(1), x.dtype)
+            for i in tf.range(10):
+                if i == x_dim:
+                    break
+                r += x[i]
+            return tf.identity(r, name="output")
+        input_names_with_port = ["input_1:0"]
+        feed_dict = {"input_1:0": np.arange(0, 15, dtype=np.int32)}
+        output_names_with_port = ["output:0"]
+        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port)
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/controlflow.py

@@ -730,7 +730,15 @@ def inline_subgraph(parent, g, scope, binding):
     for n in g.get_nodes():
         dtypes = n.output_dtypes
         shapes = n.output_shapes
-        n.graph = parent
+        subgraphs = n.get_body_graphs()
+
+        n.graph = parent # we must change node graph exactly here so that previous/following code can work
+
+        # if n has subgraphs, we need to set the correct parent graph for them
+        if subgraphs:
+            for sub_name, sub_graph in subgraphs.items():
+                n.set_body_graph_as_attr(sub_name, sub_graph)
+
         for name, shape, dtype in zip(n.output, shapes, dtypes):
             # FIXME: don't access this directly
             parent._output_shapes[name] = shape  # pylint: disable=protected-access