Function InitGraph

tensorflow/compiler/tf2xla/tf2xla.cc:362–399

InitGraph creates a graph based on the graph_def, that may then be converted to an xla::XlaComputation via ConvertGraphToXla. The graph is rewritten with _Arg and _Retval nodes, representing the inputs and outputs of the function that will be compiled. Each feed id causes a new _Arg node to be created, where we first collect all existing edges pointing from the named node's output index, and the

// with a new edge pointing from the named node's output index to that _Retval
// node.
Status InitGraph(const GraphDef& graph_def, const tf2xla::Config& config,
                 std::unique_ptr<Graph>* graph) {
  TF_RETURN_IF_ERROR(ValidateConfig(config));

  FunctionLibraryDefinition flib_def(OpRegistry::Global(), graph_def.library());
  std::unique_ptr<Graph> g(new Graph(flib_def));

  // Replace references to fed tensors with references to newly added
  // placeholders.
  GraphDef first_copy_def = graph_def;

  // Maps from name:port of a feed to the name:port of the placeholder to use.
  std::unordered_map<string, string> feed_remapping;
  TF_RETURN_IF_ERROR(AddPlaceholdersForFeeds(config, g->op_registry(),
                                             &feed_remapping, &first_copy_def));

  // Prune the GraphDef first so that unknown ops that we aren't compiling get
  // filtered out.
  GraphDef second_copy_def;
  TF_RETURN_IF_ERROR(
      PruneGraphDefInto(config, first_copy_def, &second_copy_def));

  TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(
      &second_copy_def, *g->op_registry(), /*node_offset=*/0));

  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
      GraphConstructorOptions(), std::move(second_copy_def), g.get()));
  TF_RETURN_IF_ERROR(RewriteAndPruneGraph(g.get(), config, feed_remapping));

  // Functionalize control flow.
  TF_RETURN_IF_ERROR(FunctionalizeControlFlow(g.get(), &flib_def));
  // After control flow functionalization, we might have more FunctionDef's
  // (then/else branch, loop body). Add them to the graph.
  TF_RETURN_IF_ERROR(g->AddFunctionLibrary(flib_def.ToProto()));

  *graph = std::move(g);
  return Status::OK();
}

}  // namespace