Function AppendFlatMap

tensorflow/core/grappler/optimizers/data/rebatch.cc:513–548

Rewrite graph to add `.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x). batch(minibatch_size, drop_remainder=False))` after the batch node. This ensures that the sum of the minibatch sizes in a step adds up to the global batch size. However, since this adds additional data copies (both from_tensor_slices and batch), we only use this approach when necessary, i.e. when we need to drop remai

// global batch, or when the global batch size does not divide num_replicas
// evenly.
Status AppendFlatMap(const NodeDef& batch_node, int64 num_replicas,
                     FunctionLibraryDefinition* flib, MutableGraphView* graph) {
  // `.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x).
  //     batch(minibatch_size, drop_remainder=False))`
  FunctionDef flat_map_fn;
  FunctionDefLibrary lib = flib->ToProto();
  graph_utils::SetUniqueGraphFunctionName("rebatch/flat_map_fn", &lib,
                                          &flat_map_fn);
  DataTypeVector dtypes;
  TF_RETURN_IF_ERROR(
      graph_utils::GetDatasetOutputTypesAttr(batch_node, &dtypes));
  TF_RETURN_IF_ERROR(
      CreateFlatMapFnWithBatch(dtypes, num_replicas, &flat_map_fn));

  NodeDef* flat_map_node;

  AttrValue output_shapes = batch_node.attr().at(kOutputShapesAttr);
  for (auto& shape : *output_shapes.mutable_list()->mutable_shape()) {
    if (!shape.unknown_rank() && shape.dim(0).size() != -1) {
      // Because the flat map function uses drop_remainder = False,
      // the shape might be unknown
      auto old_dim = shape.dim(0).size();
      auto new_dim = old_dim % num_replicas == 0 ? old_dim / num_replicas : -1;
      shape.mutable_dim(0)->set_size(new_dim);
    }
  }

  TF_RETURN_IF_ERROR(AddFlatMapNode(strings::StrCat(batch_node.name(), ":0"),
                                    {}, {}, flat_map_fn, output_shapes, dtypes,
                                    flib, graph, &flat_map_node));

  TF_RETURN_IF_ERROR(
      graph->UpdateFanouts(batch_node.name(), flat_map_node->name()));

  return Status::OK();
}

// There are several things we do here, depending on the values of
// batch_size and drop_remainder.