code: def flatten_model_graph_without_layerchoice(
      ir_model: GraphModelSpace,
  ):
      """
      Flatten the subgraph into root graph and jump all layerchoice
      """
  
      def _flatten_without_layerchoice(graph: Graph):
          """
          flatten this graph
          """
          model = graph.model
          node_to_remove = []
  
          for node in graph.hidden_nodes:
              if is_layerchoice_node(node):
                  for in_edge in node.incoming_edges:
                      graph.del_edge(in_edge)
                  for out_edge in node.outgoing_edges:
                      graph.del_edge(out_edge)
                  del model.graphs[node.name]
                  node.remove()
                  return
  
              node_graph = model.graphs.get(node.name)
              if node_graph is not None:
                  _flatten_without_layerchoice(node_graph)
  
                  # flatten node graph into this graph
                  id_to_new_node = {}
                  for node_graph_node in node_graph.hidden_nodes:
                      new_node = Node(
                          graph,
                          node_graph_node.id,
                          node_graph_node.name,
                          node_graph_node.operation,
                          _internal=True,
                      )
                      new_node.update_label(node_graph_node.label)
                      new_node._register()
                      id_to_new_node[new_node.id] = new_node
  
                  # reconnect node edges
                  for in_edge in node.incoming_edges:
                      graph.del_edge(in_edge)
                      for (
                          input_node_edge
                      ) in node_graph.input_node.outgoing_edges:
                          if (
                              input_node_edge.head_slot
                              == in_edge.tail_slot
                          ):
                              graph.add_edge(
                                  head=(
                                      in_edge.head,
                                      in_edge.head_slot,
                                  ),
                                  tail=(
                                      id_to_new_node[
                                          input_node_edge.tail.id
                                      ],
                                      input_node_edge.tail_slot,
                                  ),
                              )
  
                  for out_edge in node.outgoing_edges:
                      graph.del_edge(out_edge)
                      for (
                          output_node_edge
                      ) in node_graph.output_node.incoming_edges:
                          if (
                              output_node_edge.head_slot
                              == out_edge.tail_slot
                          ):
                              graph.add_edge(
                                  head=(
                                      id_to_new_node[
                                          output_node_edge.head.id
                                      ],
                                      output_node_edge.head_slot,
                                  ),
                                  tail=(
                                      out_edge.tail,
                                      out_edge.tail_slot,
                                  ),
                              )
  
                  for edge in node_graph.edges:
                      if (
                          edge.head == node_graph.input_node
                          or edge.tail == node_graph.output_node
                      ):
                          continue
                      new_head = id_to_new_node[edge.head.id]
                      new_tail = id_to_new_node[edge.tail.id]
                      Edge(
                          (new_head, edge.head_slot),
                          (new_tail, edge.tail_slot),
                          _internal=True,
                      )._register()
  
                  node_to_remove.append(node)
                  del model.graphs[node.name]
  
          for node in node_to_remove:
              node.remove()
  
      new_ir_model = ir_model.fork()
      _flatten_without_layerchoice(new_ir_model.root_graph)
  
      # remove subgraphs
      new_ir_model.graphs = {
          new_ir_model._root_graph_name: new_ir_model.root_graph
      }
      return new_ir_model
  
file_path: nni/nni/nas/space/pytorch/converter/utils.py
imports: from typing import Optional
  from typing_extensions import TypeGuard
  from nni.nas.space.graph import Cell, GraphModelSpace, Graph, Node, Edge
label_desc: Use image node_id 11 for calling a global function with example usage: flatten_model_graph_without_layerchoice(ir_model) and returns: new_ir_model