Function AddFlatMapNode

tensorflow/core/grappler/optimizers/data/rebatch.cc:401–428

}

Status AddFlatMapNode(const string& input_dataset,
                      gtl::ArraySlice<string> other_arguments,
                      gtl::ArraySlice<DataType> t_arguments,
                      const FunctionDef& flat_map_fn,
                      const AttrValue& output_shapes,
                      const DataTypeVector& output_types,
                      FunctionLibraryDefinition* flib, MutableGraphView* graph,
                      NodeDef** result) {
  TF_RETURN_IF_ERROR(flib->AddFunctionDef(flat_map_fn));
  AttrValue f;
  f.mutable_func()->set_name(flat_map_fn.signature().name());

  NodeDef flat_map_node;
  flat_map_node.set_op("FlatMapDataset");
  flat_map_node.add_input(input_dataset);
  for (const auto& arg : other_arguments) {
    flat_map_node.add_input(arg);
  }
  AddNodeAttr("f", f, &flat_map_node);
  AddNodeAttr("Targuments", t_arguments, &flat_map_node);
  AddNodeAttr(kOutputShapesAttr, output_shapes, &flat_map_node);
  AddNodeAttr(kOutputTypesAttr, output_types, &flat_map_node);

  graph_utils::SetUniqueGraphNodeName("rebatch/flat_map", graph->graph(),
                                      &flat_map_node);
  *result = graph->AddNode(std::move(flat_map_node));
  return Status::OK();
}

// def flat_map_fn(*batched_components):
//   batch_size = tf.shape(batched_components[0])[0]