add support dynamic max_iterations in loop v2

Signed-off-by: wangqiaoshi <[email protected]>

Author: wangqiaoshi

tf2onnx/onnx_opset/controlflow.py

@@ -366,15 +366,7 @@ def version_7(cls, ctx, node, **kwargs):
         # consumers, modify it in place. Otherwise, make a new const node and leave the original unchanged.
         # if maximum_iterations is not const,should add an cast node(cast to int64)
         maximum_iterations_name = node.input[1]
-        cast_mark = False
-        if node.inputs[1].type != "Const":
-            cast_mark = True
-        if cast_mark:
-            cast_inputs = [maximum_iterations_name]
-            attr = {"to": onnx_pb.TensorProto.INT64}
-            cast_name = node.name + "_cast"
-            cast_node = ctx.make_node("Cast", cast_inputs, attr, name=cast_name)
-        else:
+        if node.inputs[1].is_const():
             maximum_iterations = node.inputs[1].get_tensor_value()
             if maximum_iterations == -1:
                 maximum_iterations = np.iinfo(np.int64).max
@@ -386,6 +378,13 @@ def version_7(cls, ctx, node, **kwargs):
                 maximum_iterations_name = utils.make_name(node.inputs[1].name)
             ctx.make_const(maximum_iterations_name, np.array(maximum_iterations, dtype=np.int64))
             ctx.replace_input(node, node.input[1], maximum_iterations_name, 1)
+            maximum_iterations_int64 = maximum_iterations_name
+        else:
+            cast_inputs = [maximum_iterations_name]
+            attr = {"to": onnx_pb.TensorProto.INT64}
+            cast_name = node.name + "_cast"
+            cast_node = ctx.make_node("Cast", cast_inputs, attr, name=cast_name)
+            maximum_iterations_int64 = cast_node.output[0]
 
         cond_name = node.get_attr_str("cond")
         cond_graph = find_function(cond_name)
@@ -461,16 +460,10 @@ def version_7(cls, ctx, node, **kwargs):
         output_names = output_names[2:]
 
         branches = {"body": body}
-        if cast_mark:
-            loop_node = ctx.make_node("Loop", [cast_node.output[0], cond_outputs[0]] + loop_vars,
-                                      output_count=len(output_shapes), name=node.name + "_loop",
-                                      shapes=output_shapes, dtypes=output_dtypes, skip_conversion=True,
-                                      branches=branches)
-        else:
-            loop_node = ctx.make_node("Loop", [maximum_iterations_name, cond_outputs[0]] + loop_vars,
-                                      output_count=len(output_shapes), name=node.name + "_loop",
-                                      shapes=output_shapes, dtypes=output_dtypes, skip_conversion=True,
-                                      branches=branches)
+        loop_node = ctx.make_node("Loop", [maximum_iterations_int64, cond_outputs[0]] + loop_vars,
+                                  output_count=len(output_shapes), name=node.name + "_loop",
+                                  shapes=output_shapes, dtypes=output_dtypes, skip_conversion=True,
+                                  branches=branches)
 
         output_map = dict(zip(output_names, loop_node.output))