| 1410 | } |
| 1411 | |
| 1412 | Status SavedModelOptimizer::FindKvResourceImportNode( |
| 1413 | std::unordered_map<std::string, std::vector<Node*>>& var_parts, |
| 1414 | std::unordered_map<std::string, std::vector<Node*>>& import_nodes) { |
| 1415 | for (auto vp : var_parts) { |
| 1416 | for (Node* n : vp.second) { |
| 1417 | for (const Edge* edge : n->out_edges()) { |
| 1418 | if (edge->dst()->op_def().name() == "KvResourceImportV2") { |
| 1419 | import_nodes[vp.first].push_back(edge->dst()); |
| 1420 | } |
| 1421 | } |
| 1422 | } |
| 1423 | // NOTE: Why import_nodes num is twice to var? |
| 1424 | // Please see the details in KvResourceImport defination. |
| 1425 | if (import_nodes.find(vp.first) == import_nodes.end() || |
| 1426 | import_nodes[vp.first].size() != 2 * var_parts[vp.first].size()) { |
| 1427 | LOG(FATAL) << "KvVarHandle node count should be twice to KvResourceImportV2 node count " |
| 1428 | << ", " << vp.first << ", " << import_nodes[vp.first].size() << " VS " << var_parts[vp.first].size(); |
| 1429 | } |
| 1430 | } |
| 1431 | |
| 1432 | return Status::OK(); |
| 1433 | } |
| 1434 | |
| 1435 | Status SavedModelOptimizer::FindVariableParts( |
| 1436 | std::unordered_map<std::string, std::vector<Node*>>& var_parts) { |