| 2257 | } |
| 2258 | |
| 2259 | Status SavedModelOptimizer::GetIncrRestoreOpInputs( |
| 2260 | const Node* restore_op, |
| 2261 | std::vector<SrcInfo>& input_nodes) { |
| 2262 | for (const Edge* edge : restore_op->in_edges()) { |
| 2263 | const Node* src = edge->src(); |
| 2264 | if (src->type_string() != "NoOp") { |
| 2265 | LOG(FATAL) << "Restore op " << restore_op->name() |
| 2266 | << " has input node which type is not NoOp." |
| 2267 | << src->DebugString(); |
| 2268 | } |
| 2269 | |
| 2270 | for (const Edge* inner_edge : src->in_edges()) { |
| 2271 | Node* inner_src = inner_edge->src(); |
| 2272 | // KvResourceImportV2 -> KvResourceInsert |
| 2273 | if (inner_src->op_def().name() == "KvResourceImportV2") { |
| 2274 | if (inner_edge->src_output() != -1) { |
| 2275 | LOG(FATAL) << "Invalid tensorflow graph, restore_shard op is: " |
| 2276 | << src->DebugString() << ", KvResourceImportV2 op is: " |
| 2277 | << inner_src->DebugString(); |
| 2278 | } |
| 2279 | |
| 2280 | Node* insert_op = nullptr; |
| 2281 | Status s = ConvertKvImportToKvInsert(inner_src, &insert_op); |
| 2282 | TF_RETURN_IF_ERROR(s); |
| 2283 | |
| 2284 | input_nodes.push_back({insert_op, inner_edge->src_output()}); |
| 2285 | } else { |
| 2286 | input_nodes.push_back({inner_src, inner_edge->src_output()}); |
| 2287 | } |
| 2288 | } |
| 2289 | } |
| 2290 | |
| 2291 | return Status::OK(); |
| 2292 | } |
| 2293 | |
| 2294 | Status SavedModelOptimizer::AddIncrRestoreOps() { |
| 2295 | const std::string restore_op_name = |
nothing calls this directly
no test coverage detected