Given the constant foldable nodes in 'nodes', returns a new graph 'g'. 'g' will contain copies of the nodes in 'nodes'. In addition, if there is an edge going from a node 'n' in 'nodes' to another node in 'orig_graph' but not in 'nodes', then 'tensors_to_fetch' will contain the mapping from the corresponding copy of 'n' and the edge number in 'g' to 'n'.
| 421 | // 'nodes', then 'tensors_to_fetch' will contain the mapping from the |
| 422 | // corresponding copy of 'n' and the edge number in 'g' to 'n'. |
| 423 | Graph* GetConstantGraph( |
| 424 | const Graph* orig_graph, const std::vector<Node*>& nodes, |
| 425 | const std::unordered_map<const Node*, std::vector<Tensor>>& |
| 426 | shape_replacement_map, |
| 427 | std::map<NodeAndOutput, NodeAndOutput>* tensors_to_fetch, |
| 428 | const ConstantFoldNameGenerator& generate_new_name) { |
| 429 | Graph* constant_graph = new Graph(orig_graph->op_registry()); |
| 430 | std::unordered_map<Node*, std::vector<Node*>> node_map; |
| 431 | node_map[orig_graph->source_node()] = {constant_graph->source_node()}; |
| 432 | node_map[orig_graph->sink_node()] = {constant_graph->sink_node()}; |
| 433 | for (Node* n : nodes) { |
| 434 | if (shape_replacement_map.count(n) == 0) { |
| 435 | AddNodeToConstantGraph(n, &node_map, constant_graph); |
| 436 | } else { |
| 437 | AddShapeNodeToConstantGraph(n, shape_replacement_map, &node_map, |
| 438 | generate_new_name, constant_graph); |
| 439 | } |
| 440 | } |
| 441 | |
| 442 | for (auto const& added_nodes : node_map) { |
| 443 | for (const Edge* out_edge : added_nodes.first->out_edges()) { |
| 444 | if (node_map.count(out_edge->dst()) == 0) { |
| 445 | if (out_edge->IsControlEdge()) continue; |
| 446 | if (added_nodes.second.size() == 1) { |
| 447 | tensors_to_fetch->insert( |
| 448 | {{added_nodes.second[0], out_edge->src_output()}, |
| 449 | {added_nodes.first, out_edge->src_output()}}); |
| 450 | } else { |
| 451 | // The node had multiple outputs and was replaced by a |
| 452 | // vector of constants, so the NodeAndOutput is the 0th |
| 453 | // output of the kth added constant, rather than the kth |
| 454 | // output of the added node as in the standard case above. |
| 455 | tensors_to_fetch->insert( |
| 456 | {{added_nodes.second[out_edge->src_output()], 0}, |
| 457 | {added_nodes.first, out_edge->src_output()}}); |
| 458 | } |
| 459 | } |
| 460 | } |
| 461 | } |
| 462 | |
| 463 | return constant_graph; |
| 464 | } |
| 465 | |
| 466 | // Replaces the identified Tensor in 'graph' by a 'Const' node with |
| 467 | // the value supplied in 'constant'. 'partition_device', if non-null |
no test coverage detected