| 121 | } |
| 122 | |
| 123 | Status FetchOutputs( |
| 124 | Graph* g, const std::vector<std::unique_ptr<PruneRewrite>>& fetch_rewrites, |
| 125 | NameIndex* name_index, std::vector<Node*>* out_fetch_nodes, |
| 126 | DataTypeVector* out_fetch_types) { |
| 127 | out_fetch_nodes->clear(); |
| 128 | out_fetch_nodes->reserve(fetch_rewrites.size()); |
| 129 | for (size_t i = 0; i < fetch_rewrites.size(); ++i) { |
| 130 | const string& t = fetch_rewrites[i]->endpoint_name(); |
| 131 | |
| 132 | // Parse t into node_name and output_index. |
| 133 | TensorId id(ParseTensorName(t)); |
| 134 | |
| 135 | // Find node in graph with that name. |
| 136 | auto iter = name_index->find(id.first); |
| 137 | if (iter == name_index->end()) { |
| 138 | return errors::NotFound("FetchOutputs node ", t, ": not found"); |
| 139 | } |
| 140 | Node* n = iter->second; |
| 141 | DCHECK_EQ(n->name(), id.first); |
| 142 | VLOG(2) << "Found fetch node for " << t; |
| 143 | |
| 144 | // Validate output_index |
| 145 | if (n->num_outputs() == 0) { |
| 146 | return errors::InvalidArgument( |
| 147 | "Tried to fetch data for '", t, |
| 148 | "', which produces no output. To run to a node but not fetch any " |
| 149 | "data, pass '", |
| 150 | t, |
| 151 | "' as an argument to the 'target_node_names' argument of the " |
| 152 | "Session::Run API."); |
| 153 | } else if (id.second >= n->num_outputs()) { |
| 154 | return errors::InvalidArgument("FetchOutputs ", t, |
| 155 | ": output index too large, must be < ", |
| 156 | n->num_outputs()); |
| 157 | } |
| 158 | |
| 159 | // Create the fetch Node and connect it up |
| 160 | Node* fetch_node; |
| 161 | TF_RETURN_IF_ERROR( |
| 162 | fetch_rewrites[i]->AddNode(g, {n, id.second}, &fetch_node)); |
| 163 | |
| 164 | // Update the index. |
| 165 | (*name_index)[fetch_node->name()] = fetch_node; |
| 166 | |
| 167 | // Duplicate control edges aren't allowed, but fetch_node was *just* created |
| 168 | // so there's no need to check for a duplicate. |
| 169 | g->AddControlEdge(fetch_node, g->sink_node(), true); |
| 170 | out_fetch_nodes->push_back(fetch_node); |
| 171 | out_fetch_types->push_back(BaseType(n->output_type(id.second))); |
| 172 | } |
| 173 | |
| 174 | return Status::OK(); |
| 175 | } |
| 176 | |
| 177 | bool AddNodeToTargets(const string& node_or_tensor_name, |
| 178 | const NameIndex& name_index, |
no test coverage detected