Function Partition

tensorflow/core/graph/graph_partition_test.cc:86–118

}

void Partition(const GraphDef& graph_def,
               std::unordered_map<string, GraphDef>* partitions) {
  Graph g(OpRegistry::Global());
  GraphConstructorOptions opts;
  TF_CHECK_OK(ConvertGraphDefToGraph(opts, graph_def, &g));

  // Assigns devices to each node. Uses 1st letter of the node name as the
  // device index if no device is specified.
  for (Node* node : g.nodes()) {
    string device_name = !node->requested_device().empty()
                             ? node->requested_device()
                             : DeviceName(node);
    node->set_assigned_device_name(device_name);
  }

  PartitionOptions popts;
  popts.node_to_loc = SplitByDevice;
  popts.new_name = [&g](const string& prefix) { return g.NewName(prefix); };
  popts.get_incarnation = [](const string& name) {
    return (name[0] - 'A') + 100;
  };
  Status s = Partition(popts, &g, partitions);
  CHECK(s.ok()) << s;

  // Check versions.
  EXPECT_EQ(graph_def.versions().producer(), TF_GRAPH_DEF_VERSION);
  // Partitions must inherit the versions of the original graph.
  for (auto& it : *partitions) {
    EXPECT_EQ(graph_def.versions().producer(), it.second.versions().producer());
    EXPECT_EQ(graph_def.versions().min_consumer(),
              it.second.versions().min_consumer());
  }
}

void CheckLoopConstruction(const GraphDef& graph_def) {
  std::unordered_map<string, GraphDef> partitions;