torch/csrc/jit/interpreter.cpp
// removes all inputs and outputs to a graph, replacing them with nodes before of after each insertStage
removes all inputs and outputs to a graph, replacing them with nodes before of after each insertStage
removes all inputs and outputs to a graph replacing them with nodes before of after each insertStage
static std::vector<std::vector<TypePtr>> flattenStages(Graph & graph) {
  std::vector<std::vector<TypePtr>> stage_input_types;
  WithInsertPoint guard(*graph.nodes().begin());
  size_t input_pos = 0;
  size_t output_pos = 0;
  auto it = graph.nodes().begin();
  for(size_t i = 0; i <= graph.stage(); i++) {
    stage_input_types.emplace_back();
    auto store = graph.create(prim::Store, 0)->insertBefore(*it);
    while(input_pos < graph.inputs().size() && graph.inputs()[input_pos]->stage() == i) {
      auto nv = store->addOutput();
      auto old_node = graph.inputs()[input_pos];
      nv->setType(old_node->type());
      stage_input_types[i].push_back(old_node->type());
      old_node->replaceAllUsesWith(nv);
      input_pos++;
    }
    while(it != graph.nodes().end() && it->stage() == i)
      ++it;
    auto load = graph.create(prim::Load, 0)->insertBefore(*it);
    while(output_pos < graph.outputs().size() && graph.outputs()[output_pos]->stage() == i) {
      load->addInput(graph.outputs()[output_pos]);
      output_pos++;
    }
  }
  while (graph.inputs().size() > 0)
    graph.eraseInput(graph.inputs().size() - 1);
  while (graph.outputs().size() > 0)
    graph.eraseOutput(graph.outputs().size() - 1);
  return stage_input_types;
}
static std :: vector < std :: vector < TypePtr > > flattenStages ( Graph & graph ) { std :: vector < std :: vector < TypePtr > > stage_input_types ; WithInsertPoint guard ( * graph . nodes ( ) . begin ( ) ) ; size_t input_pos = 0 ; size_t output_pos = 0 ; auto it = graph . nodes ( ) . begin ( ) ; for ( size_t i = 0 ; i <= graph . stage ( ) ; i ++ ) { stage_input_types . emplace_back ( ) ; auto store = graph . create ( prim :: Store , 0 ) -> insertBefore ( * it ) ; while ( input_pos < graph . inputs ( ) . size ( ) && graph . inputs ( ) [ input_pos ] -> stage ( ) == i ) { auto nv = store -> addOutput ( ) ; auto old_node = graph . inputs ( ) [ input_pos ] ; nv -> setType ( old_node -> type ( ) ) ; stage_input_types [ i ] . push_back ( old_node -> type ( ) ) ; old_node -> replaceAllUsesWith ( nv ) ; input_pos ++ ; } while ( it != graph . nodes ( ) . end ( ) && it -> stage ( ) == i ) ++ it ; auto load = graph . create ( prim :: Load , 0 ) -> insertBefore ( * it ) ; while ( output_pos < graph . outputs ( ) . size ( ) && graph . outputs ( ) [ output_pos ] -> stage ( ) == i ) { load -> addInput ( graph . outputs ( ) [ output_pos ] ) ; output_pos ++ ; } } while ( graph . inputs ( ) . size ( ) > 0 ) graph . eraseInput ( graph . inputs ( ) . size ( ) - 1 ) ; while ( graph . outputs ( ) . size ( ) > 0 ) graph . eraseOutput ( graph . outputs ( ) . size ( ) - 1 ) ; return stage_input_types ; }
removes all inputs and outputs to a graph, replacing them with nodes before of after each insertStage
removes all inputs and outputs to a graph replacing them with nodes before of after each insertStage
// because JIT classic needs this to fix up gradients, remove when possible