Rewrite graph by replacing the output tensors specified in "fed_outputs" with special feed nodes for each specified output tensor, and removing any nodes that are now disconnected from the part of the graph that reaches the sink node. The set of special feed nodes added to the graph are returned in "*feed_nodes". Return true on success. On error, return false and sets *error to an appropriate e
| 56 | // an appropriate error message (and *g is left in an indeterminate |
| 57 | // state). |
| 58 | Status FeedInputs( |
| 59 | Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& feed_rewrites, |
| 60 | NameIndex* name_index, DataTypeVector* out_feed_types) { |
| 61 | out_feed_types->clear(); |
| 62 | out_feed_types->reserve(feed_rewrites.size()); |
| 63 | for (size_t i = 0; i < feed_rewrites.size(); ++i) { |
| 64 | const string& t = feed_rewrites[i]->endpoint_name(); |
| 65 | TensorId id(ParseTensorName(t)); |
| 66 | |
| 67 | auto iter = name_index->find(id.first); |
| 68 | if (iter == name_index->end()) { |
| 69 | return errors::NotFound("FeedInputs: unable to find feed output ", t); |
| 70 | } |
| 71 | Node* n = iter->second; |
| 72 | DCHECK_EQ(n->name(), id.first); |
| 73 | if (id.second >= n->num_outputs()) { |
| 74 | return errors::InvalidArgument( |
| 75 | "FeedInputs: ", t, " should have output index < ", n->num_outputs()); |
| 76 | } |
| 77 | |
| 78 | Node* feed_node; |
| 79 | TF_RETURN_IF_ERROR( |
| 80 | feed_rewrites[i]->AddNode(g, {n, id.second}, &feed_node)); |
| 81 | |
| 82 | // Update name_index |
| 83 | (*name_index)[feed_node->name()] = feed_node; |
| 84 | // Duplicate control edges aren't allowed, but feed_node was *just* created |
| 85 | // so there's no need to check for a duplicate. |
| 86 | g->AddControlEdge(g->source_node(), feed_node, true); |
| 87 | |
| 88 | // Look through edges coming out of "n" for edges whose src_output() index |
| 89 | // matches "output_index". If found, replace the edges with a connection |
| 90 | // from the special feed node. |
| 91 | std::vector<const Edge*> to_remove; |
| 92 | for (const Edge* e : n->out_edges()) { |
| 93 | if (e->src_output() == id.second) { |
| 94 | to_remove.emplace_back(e); |
| 95 | } else if (e->src_output() == Graph::kControlSlot && |
| 96 | (n->type_string() == "Placeholder" || |
| 97 | n->type_string() == "PlaceholderV2")) { |
| 98 | // When feeding a Placeholder node, any outgoing control edges |
| 99 | // will be replaced with a control edge from the replacement |
| 100 | // feed_node. |
| 101 | // TODO(josh11b,mrry): Come up with a more elegant way of addressing |
| 102 | // the general version of this problem. |
| 103 | to_remove.emplace_back(e); |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | for (const Edge* e : to_remove) { |
| 108 | if (e->src_output() == id.second) { |
| 109 | g->AddEdge(feed_node, 0, e->dst(), e->dst_input()); |
| 110 | } else { |
| 111 | CHECK_EQ(Graph::kControlSlot, e->src_output()); |
| 112 | // Duplicate control edges aren't allowed, but feed_node was *just* |
| 113 | // created so there's no need to check for a duplicate. |
| 114 | g->AddControlEdge(feed_node, e->dst(), true); |
| 115 | } |
no test coverage detected