Function OptimizeGraph

tensorflow/core/grappler/optimizers/data/rebatch.cc:833–862

Helper function that given a GrapplerItem generates a mutated graph def with the batch size changed. The GrapplerItem could be generated from the main graph or could be a function graph.

// with the batch size changed. The GrapplerItem could be generated from the
// main graph or could be a function graph.
Status OptimizeGraph(const GrapplerItem& item, int64 num_replicas,
                     bool use_fallback, GraphDef* output) {
  *output = item.graph;
  MutableGraphView graph(output);

  FunctionLibraryDefinition flib(OpRegistry::Global(), item.graph.library());

  NodeDef* sink_node;
  TF_RETURN_IF_ERROR(graph_utils::GetFetchNode(graph, item, &sink_node));

  Status s = RecursivelyHandleOp(*sink_node, num_replicas, use_fallback, &flib,
                                 &graph);
  if (!s.ok()) {
    if (use_fallback) {
      VLOG(1) << "Failed to rebatch by rewriting the batch transformation ("
              << s << "). Using a fallback method instead.";
      // If RecursivelyHandleOp fails, we reset `graph` to use the original,
      // graph, since that function may have mutated `graph`.
      *output = item.graph;
      graph = MutableGraphView(output);
      TF_RETURN_IF_ERROR(
          RebatchWithFallback(sink_node, num_replicas, &flib, &graph));
    } else {
      // Return the error
      return s;
    }
  }
  *output->mutable_library() = flib.ToProto();
  return Status::OK();
}

}  // anonymous namespace