code: def _rewrite_cond_functions(tf_ssa, fn):
      r"""
      Rewrite tf.cond'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 'StatelessIf' and 'If'.
  
      Parameters
      ----------
      tf_ssa: NetworkEnsemble
          An object that contains multiple functions / sub-graphs.
      fn: SSAFunction
          Function that contains graph to operate on.
  
      Examples
      --------
  
      Input:
  
          Before pass "main" graph:
  
              [const/greater/y] ---------\
              [placeholder/args_0] -> [greater] -> [if] -> [identity]
                                \------------------/  \--> [identity]
              [placeholder/args_1] ----------------/
  
          Before pass "then" graph:
  
              [const/sub/y] ---------------\
              [placeholder/sub_args_0] -> [sub]
              [placeholder/sub_args_1] -> [identity]
  
          Before pass "else" graph:
  
              [const/add/y] ---------------\
              [placeholder/add_args_0] -> [add]
  
              [const/mul/y] ---------------\
              [placeholder/add_args_1] -> [mul]
  
      Output:
  
          After pass "main" graph:
  
              [const/greater/y] ---------\
              [placeholder/args_0] -> [greater] -> [make_tuple] -> [if] -> [get_tuple] -> [identity]
                                \---------------------/               \--> [get_tuple] -> [identity]
              [placeholder/args_1] -------------------/
  
          After pass "then" graph:
  
                                        [const/sub/y] ---------------\
              [entry] -> [get_tuple] -> [placeholder/sub_args_0] -> [sub] -> [make_tuple] -> [return]
                      -> [get_tuple] -> [placeholder/sub_args_1] -----------------/
  
          After pass "else" graph:
  
                                        [const/add/y] ---------------\
              [entry] -> [get_tuple] -> [placeholder/add_args_0] -> [add] -> [make_tuple] -> [return]
                      -> [get_tuple] -> [placeholder/add_args_1] -> [mul] --------/
                                        [const/mul/y] ---------------/
  
      """
      for cond_name, cond_node in fn.graph.copy().items():
          if cond_node.op not in {"StatelessIf", "If"}:
              continue
  
          then_fn_name = cond_node.attr.get("then_branch")
          else_fn_name = cond_node.attr.get("else_branch")
  
          msg = "Rewriting '{}' ({}) sub-graphs: then '{}', else '{}'"
          logging.info(
              msg.format(
                  cond_node.name,
                  cond_node.op,
                  then_fn_name,
                  else_fn_name,
              )
          )
  
          then_fn = tf_ssa.functions.get(then_fn_name)
          else_fn = tf_ssa.functions.get(else_fn_name)
  
          # insert function entry nodes
          then_entry = _insert_function_entry(then_fn)
          else_entry = _insert_function_entry(else_fn)
  
          # pack node inputs to a single tuple
          cond_input = _insert_make_tuple(
              fn, "make_tuple/{}".format(cond_name)
          )
          for ci in cond_node.inputs:
              disconnect_edge(fn.graph, ci, cond_node.name)
              connect_edge(fn.graph, ci, cond_input)
          connect_edge(fn.graph, cond_input, cond_node.name)
  
          # unpack node outputs to multiple get_tuples
          for i, co in enumerate(cond_node.outputs):
              # utilize FunctionDef's ret to make sure function outputs and
              # node outputs order matches when multiple outputs are there.
              # Fallback to use original cond_node.outputs order if fails.
              o_original = fn.graph[co].original_node
              if o_original:
                  c_input = [
                      n
                      for n in o_original.input
                      if str(n).startswith(cond_name)
                  ][0]
                  c_index = (
                      c_input.split(":")[-1] if ":" in c_input else 0
                  )
                  mapped_name = then_fn.ret[
                      "identity_{}".format(c_index)
                  ].split(":")[0]
                  if mapped_name in then_fn.outputs:
                      idx = then_fn.outputs.index(mapped_name)
                  else:  # in else_fn.outputs
                      idx = else_fn.outputs.index(mapped_name)
              else:
                  idx = i
  
              cond_output = _insert_get_tuple(
                  fn, "get_tuple/{}/{}".format(idx, cond_name), idx
              )
              edge_idx = fn.graph[co].inputs.index(cond_node.name)
              replace_dest(fn.graph, cond_node, co, cond_output)
              connect_edge_at_index(fn.graph, cond_output, co, edge_idx)
  
          # fetch inputs using get_tuple for then branch
          for i, ti in enumerate(then_fn.inputs):
              then_input = _insert_get_tuple(
                  then_fn, "get_tuple/{}/{}".format(i, ti), i + 1
              )
              connect_edge(then_fn.graph, then_entry, then_input)
              replace_node(then_fn.graph, ti, then_input)
              delete_node(then_fn.graph, ti)
  
          # fetch inputs using get_tuple for else branch
          for i, ei in enumerate(else_fn.inputs):
              else_input = _insert_get_tuple(
                  else_fn, "get_tuple/{}/{}".format(i, ei), i + 1
              )
              connect_edge(else_fn.graph, else_entry, else_input)
              replace_node(else_fn.graph, ei, else_input)
              delete_node(else_fn.graph, ei)
  
          # returns a tuple of value(s) as output for then branch
          then_output = _insert_make_tuple(then_fn)
          for to in then_fn.outputs:
              if to not in then_fn.graph.keys():
                  # from identity, map back to get_tuple node
                  to = "get_tuple/{}/{}".format(
                      then_fn.inputs.index(to), to
                  )
              connect_edge(then_fn.graph, to, then_output.name)
  
          then_return = _insert_return(then_fn)
          connect_edge(
              then_fn.graph, then_output.name, then_return.name
          )
  
          # returns a tuple of value(s) as output for else branch
          else_output = _insert_make_tuple(else_fn)
          for eo in else_fn.outputs:
              if eo not in else_fn.graph.keys():
                  # from identity, map back to get_tuple node
                  eo = "get_tuple/{}/{}".format(
                      else_fn.inputs.index(eo), eo
                  )
              connect_edge(else_fn.graph, eo, else_output.name)
  
          else_return = _insert_return(else_fn)
          connect_edge(
              else_fn.graph, else_output.name, else_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 8 for calling a global function with example usage: _rewrite_cond_functions(tf_ssa, fn) without return types