Function RewriteBatchNode

tensorflow/core/grappler/optimizers/data/rebatch.cc:586–593

}

Status RewriteBatchNode(const NodeDef& batch_node, int64 num_replicas,
                        FunctionLibraryDefinition* flib,
                        MutableGraphView* graph) {
  if (ShouldMutateBatchSizeDirectly(batch_node, num_replicas, graph)) {
    return MutateBatchSize(batch_node, num_replicas, graph);
  }
  return AppendFlatMap(batch_node, num_replicas, flib, graph);
}

Status OptimizeGraph(const GrapplerItem& item, int64 num_replicas,
                     bool use_fallback, GraphDef* output);