# Buggy code
def _remove_output_identity_nodes(nnssa):
    """
    remove identity nodes that ARE connected to the model outputs
    """
    delete_count = 0
    for fn_key in list(nnssa.functions.keys()):
        f = nnssa.functions[fn_key]
        keys = list(f.graph.keys())
        for k in keys:
            if k not in f.graph:
                continue
            if f.graph[k].op == "Identity" and len(f.graph[k].inputs) == 1:
                if (
                    len(f.graph[k].outputs) == 0
                    and (k in f.outputs)
                    and k == f.graph[k].name
                ):
                    # this means node k is the "output-identity" node that nnssa inserts
                    # we remove it here
                    delete_count += 1
                    parent_name = f.graph[k].inputs[0]
                    f.graph[parent_name].control_outputs = f.graph[k].control_outputs

                    if any(
                        [an_output != k for an_output in f.graph[parent_name].outputs]
                    ):
                        continue

                    del f.graph[k]
                    f.graph[k] = copy.deepcopy(f.graph[parent_name])
                    del f.graph[parent_name]
                    f.graph[k].name = k
                    f.graph[k].outputs = []

                    for p in f.graph[k].inputs:
                        for idx, out in enumerate(f.graph[p].outputs):
                            if out == parent_name:
                                f.graph[p].outputs[idx] = k

                    for p in f.graph[k].control_inputs:
                        for idx, out in enumerate(f.graph[p].control_outputs):
                            if out == parent_name:
                                f.graph[p].control_outputs[idx] = k
    return delete_count


# Traceback (AssertionError)
0 assert nodes deleted
5 nodes deleted
0 nodes deleted
0 nodes deleted
[Op Fusion] fuse_bias_add() deleted 4 nodes.
1 identity nodes deleted
2 disconnected nodes deleted
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
~/.pyenv/versions/3.6.8/envs/fritz-models-private/src/coremltools/coremltools/converters/tensorflow/_tf_converter.py in convert(filename, inputs, outputs, image_input_names, is_bgr, red_bias, green_bias, blue_bias, gray_bias, image_scale, class_labels, predicted_feature_name, predicted_probabilities_output, add_custom_layers, custom_conversion_functions, custom_shape_functions, **kwargs)
94                                  custom_shape_functions=custom_shape_functions,
---> 95                                  optional_inputs=optional_inputs)
96     except ImportError as e:

~/.pyenv/versions/3.6.8/envs/fritz-models-private/src/coremltools/coremltools/converters/nnssa/coreml/ssa_converter.py in ssa_convert(ssa, top_func, inputs, outputs, image_input_names, is_bgr, red_bias, green_bias, blue_bias, gray_bias, image_scale, class_labels, predicted_feature_name, predicted_probabilities_output, add_custom_layers, custom_conversion_functions, custom_shape_functions, optional_inputs)
132     for f in list(ssa.functions.values()):
--> 133         check_connections(f.graph)
134

~/.pyenv/versions/3.6.8/envs/fritz-models-private/src/coremltools/coremltools/converters/nnssa/commons/basic_graph_ops.py in check_connections(gd)
152             try:
--> 153                 assert (k in gd[i].control_inputs)
154             except Exception:

# Fixed code
def _remove_output_identity_nodes(nnssa):
    """
    remove identity nodes that ARE connected to the model outputs
    """
    delete_count = 0
    for fn_key in list(nnssa.functions.keys()):
        f = nnssa.functions[fn_key]
        keys = list(f.graph.keys())
        for k in keys:
            if k not in f.graph:
                continue
            node = f.graph[k]

            if node.op != "Identity" or len(node.inputs) != 1:
                continue

            if len(node.outputs) != 0 or (k not in f.outputs) or k != node.name:
                continue
            # this means node k is the "output-identity" node that nnssa inserts
            # we remove it here
            parent_name = node.inputs[0]
            parent_node = f.graph[parent_name]

            # Continue if parent node has an other outputs than identity node.
            if any([an_output != k for an_output in parent_node.outputs]):
                continue

            delete_count += 1

            # Remove Identity node and copy existing parent node
            parent_node = copy.deepcopy(f.graph[parent_name])
            for control_input_name in node.control_inputs:
                if control_input_name == parent_node.name:
                    continue
                if control_input_name in parent_node.control_inputs:
                    continue

                parent_node.control_inputs.append(control_input_name)

            del f.graph[k]
            del f.graph[parent_name]
            parent_node.name = k
            parent_node.outputs = []
            f.graph[k] = parent_node

            node = f.graph[k]
            for p in node.inputs:
                for idx, out in enumerate(f.graph[p].outputs):
                    if out == parent_name:
                        f.graph[p].outputs[idx] = k

            for p in node.control_inputs:
                for idx, out in enumerate(f.graph[p].control_outputs):
                    if out == parent_name:
                        f.graph[p].control_outputs[idx] = k

    return delete_count