code: def _rewrite_while_loop_functions(tf_ssa, fn):
      """
      Rewrite tf.while_loop's sub-graphs with get_tuple, make_tuple,
      function_entry and return ops. This rewrite is required in order to convert
      functional form control flow v2 nodes 'StatelessWhile' and 'While'.
  
      Parameters
      ----------
      tf_ssa: NetworkEnsemble
          An object that contains multiple functions / sub-graphs.
      fn: SSAFunction
          Function that contains graph to operate on.
  
      Example
      -------
  
      Input:
  
          Before pass "main" graph:
  
              [placeholder/args_0] --> [while] --> [identity]
  
          Before pass "cond" graph:
  
              [const/mean] -------\
              [placeholder] --> [mean] --> [greater]
              [const/greater/y] --------------/
  
          Before pass "body" graph:
  
              [const/sub/y] ------\
              [placeholder] ---> [sub]
  
      Output:
  
          After pass "main" graph:
  
              [placeholder/args_0] --> [make_tuple] --> [while] --> [get_tuple] --> [identity]
  
          After pass "cond" graph:
  
                                        [const/mean] ------\
              [entry] -> [get_tuple] -> [placeholder] -> [mean] -> [greater] -> [make_tuple] -> [return]
                                        [const/greater/y] ------------/
  
          After pass "body" graph:
  
                                        [const/sub/y] ----\
              [entry] -> [get_tuple] -> [placeholder] -> [sub] -> [make_tuple] -> [return]
      """
      for while_name, while_node in fn.graph.copy().items():
          if while_node.op not in {"StatelessWhile", "While"}:
              continue
  
          cond_fn_name = while_node.attr.get("cond")
          body_fn_name = while_node.attr.get("body")
  
          msg = "Rewriting '{}' ({}) sub-graphs: cond '{}', body '{}'"
          logging.info(
              msg.format(
                  while_node.name,
                  while_node.op,
                  cond_fn_name,
                  body_fn_name,
              )
          )
  
          cond_fn = tf_ssa.functions.get(cond_fn_name)
          body_fn = tf_ssa.functions.get(body_fn_name)
  
          # insert function entry nodes
          cond_entry = _insert_function_entry(cond_fn)
          body_entry = _insert_function_entry(body_fn)
  
          # pack node inputs to a single tuple
          while_input_tuple = _insert_make_tuple(
              fn, "make_tuple/{}".format(while_name)
          )
          for wi in while_node.inputs:
              disconnect_edge(fn.graph, wi, while_node.name)
              connect_edge(fn.graph, wi, while_input_tuple)
          connect_edge(fn.graph, while_input_tuple, while_node.name)
  
          # unpack node outputs to multiple get_tuples
          for i, wo in enumerate(while_node.outputs):
              # utilize FunctionDef's ret to make sure function outputs and
              # node outputs order matches when multiple outputs are there.
              o_original = fn.graph[wo].original_node
              while_input = [
                  n
                  for n in o_original.input
                  if str(n).startswith(while_name)
              ][0]
              while_index = while_input.split(":")[-1]
              mapped_name = body_fn.ret[
                  "identity_{}".format(while_index)
              ].split(":")[0]
              idx = body_fn.outputs.index(mapped_name)
  
              loop_output = _insert_get_tuple(
                  fn, "get_tuple/{}/{}".format(idx, while_input), idx
              )
  
              edge_idx = fn.graph[wo].inputs.index(while_node.name)
              replace_dest(fn.graph, while_node, wo, loop_output)
              connect_edge_at_index(fn.graph, loop_output, wo, edge_idx)
  
          # fetch inputs using get_tuple for cond fn
          for i, ci in enumerate(cond_fn.inputs):
              cond_input = _insert_get_tuple(
                  cond_fn, "get_tuple/{}/{}".format(i, ci), i
              )
              connect_edge(cond_fn.graph, cond_entry, cond_input)
              replace_node(cond_fn.graph, ci, cond_input)
              delete_node(cond_fn.graph, ci)
  
          # fetch inputs using get_tuple for body fn
          for i, bi in enumerate(body_fn.inputs):
              new_name = "get_tuple/{}/{}".format(i, bi)
  
              if bi in body_fn.outputs:  # input is also an output
                  body_fn.outputs[body_fn.outputs.index(bi)] = new_name
  
              body_input = _insert_get_tuple(body_fn, new_name, i)
  
              connect_edge(body_fn.graph, body_entry, body_input)
              replace_node(body_fn.graph, bi, body_input)
              delete_node(body_fn.graph, bi)
  
          # returns a tuple of value(s) as output for cond fn
          cond_output = _insert_make_tuple(cond_fn)
          for co in cond_fn.outputs:
              connect_edge(cond_fn.graph, co, cond_output.name)
  
          cond_return = _insert_return(cond_fn)
          connect_edge(
              cond_fn.graph, cond_output.name, cond_return.name
          )
  
          # returns a tuple of value(s) as output for body branch
          body_output = _insert_make_tuple(body_fn)
  
          for bo in body_fn.outputs:
              connect_edge(body_fn.graph, bo, body_output.name)
  
          body_return = _insert_return(body_fn)
          connect_edge(
              body_fn.graph, body_output.name, body_return.name
          )
  
file_path: turicreate/src/external/coremltools_wrap/coremltools/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import absolute_import
  import logging
  from coremltools.converters.mil.frontend.tensorflow.parsed_tf_node import ParsedTFNode
  from coremltools.converters.mil.frontend.tensorflow.basic_graph_ops import disconnect_edge, connect_edge, delete_node, replace_node, replace_dest, connect_edge_at_index
label_desc: Use image node_id 10 for calling a global function with example usage: _rewrite_while_loop_functions(tf_ssa, fn) without return types