def _eliminate_loop_cond_nodes(tf_ssa, fn):
    """
    Eliminate loop condition nodes, such as loop_counters, max_iterations from
    the cond sub-graph and body sub-graph of tf.while_loop.

    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:

            [while/maximum_iterations] -----\
            [while/loop_counter] -------> [while] --> [identity]
            [placeholder/args_0] ----------/

        Before pass "cond" graph:

            [const/mean] -------\
            [placeholder] --> [mean] --> [greater]
            [const/greater/y] --------------/

            [while_maximum_iterations], [while_loop_counter] (not connected)

        Before pass "body" graph:

            [const/sub/y] ------\
            [placeholder] ---> [sub]

            [const/add/y] ------------\
            [while_loop_counter] --> [add]

            [while_maximum_iterations] (not connected)

    Output:

        After pass "main" graph:

            [placeholder/args_0] --> [while] --> [identity]

        After pass "cond" graph:

            [const/mean] -------\
            [placeholder] --> [mean] --> [greater]
            [const/greater/y] --------------/

        After pass "body" graph:

            [const/sub/y] ------\
            [placeholder] ---> [sub]
    """
    for name, node in fn.graph.copy().items():
        if node.op not in {"StatelessWhile", "While"}:
            continue

        cond_fn = tf_ssa.functions.get(node.attr.get("cond"))
        body_fn = tf_ssa.functions.get(node.attr.get("body"))

        cond_lc_nodes = {cond_fn.inputs.pop(0), cond_fn.inputs.pop(0)}
        logger.info("Removing {} from cond graph".format(cond_lc_nodes))
        for n in cond_lc_nodes:
            delete_node(cond_fn.graph, n)

        body_lc_nodes = {body_fn.inputs.pop(0), body_fn.inputs.pop(0)}
        q = list(body_lc_nodes)

        # delete entire sub-fn
        while len(q) > 0:
            n = body_fn.graph[q.pop(0)]
            for o in n.outputs:
                if o not in body_lc_nodes:
                    q.append(o)
                body_lc_nodes.add(o)
                for i in body_fn.graph[o].inputs:
                    if i not in body_lc_nodes:
                        q.append(i)
                    body_lc_nodes.add(i)

        # remove if in outputs
        for n in body_lc_nodes:
            if n in body_fn.outputs:
                msg = "Removing '{}' ({}) from body fn outputs"
                logger.info(msg.format(n, body_fn.graph[n].op))
                body_fn.outputs.remove(n)

        logger.info("Removing {} from body graph".format(body_lc_nodes))
        for n in body_lc_nodes:
            delete_node(body_fn.graph, n)

    Eliminate loop condition nodes, such as loop_counters, max_iterations from
    the cond sub-graph and body sub-graph of tf.while_loop.

    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:

            [while/maximum_iterations] -----            [while/loop_counter] -------> [while] --> [identity]
            [placeholder/args_0] ----------/

        Before pass "cond" graph:

            [const/mean] -------            [placeholder] --> [mean] --> [greater]
            [const/greater/y] --------------/

            [while_maximum_iterations], [while_loop_counter] (not connected)

        Before pass "body" graph:

            [const/sub/y] ------            [placeholder] ---> [sub]

            [const/add/y] ------------            [while_loop_counter] --> [add]

            [while_maximum_iterations] (not connected)

    Output:

        After pass "main" graph:

            [placeholder/args_0] --> [while] --> [identity]

        After pass "cond" graph:

            [const/mean] -------            [placeholder] --> [mean] --> [greater]
            [const/greater/y] --------------/

        After pass "body" graph:

            [const/sub/y] ------            [placeholder] ---> [sub]
    
coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py