| 1472 | } |
| 1473 | |
| 1474 | Status SavedModelOptimizer::ConvertKVOps() { |
| 1475 | |
| 1476 | // Find sparse lookup/Import ops and replace them |
| 1477 | // with KvLookup and KvImport |
| 1478 | for (Node* node : graph_.nodes()) { |
| 1479 | // Get input edges |
| 1480 | std::vector<SrcInfo> input_info; |
| 1481 | Status s_input_info = GetInputNodesInfo(&input_info, node); |
| 1482 | TF_RETURN_IF_ERROR(s_input_info); |
| 1483 | int edge_count = node->num_inputs(); |
| 1484 | |
| 1485 | std::unordered_map<std::string, const AttrValue*> attr_info; |
| 1486 | |
| 1487 | if (node->op_def().name() == "KvResourceGather") { |
| 1488 | if (!storage_pointer_node_ || !version_node_) { |
| 1489 | return tensorflow::errors::Internal( |
| 1490 | "Not found a storage pointer or version node in the graph."); |
| 1491 | } |
| 1492 | |
| 1493 | std::vector<SrcInfo> gather_input_info; |
| 1494 | // indices |
| 1495 | gather_input_info.push_back( |
| 1496 | SrcInfo{input_info[gather_input_indice_slot].src_node, |
| 1497 | input_info[gather_input_indice_slot].src_slot}); |
| 1498 | // default_value |
| 1499 | gather_input_info.push_back( |
| 1500 | SrcInfo{input_info[gather_input_default_val_slot].src_node, |
| 1501 | input_info[gather_input_default_val_slot].src_slot}); |
| 1502 | // storage pointer |
| 1503 | gather_input_info.push_back(SrcInfo{storage_pointer_node_, 0}); |
| 1504 | |
| 1505 | // model version |
| 1506 | gather_input_info.push_back(SrcInfo{version_node_, 0}); |
| 1507 | |
| 1508 | // control edges |
| 1509 | for (size_t i = edge_count; i < input_info.size(); ++i) { |
| 1510 | gather_input_info.push_back(input_info[i]); |
| 1511 | } |
| 1512 | |
| 1513 | AttrValue feature_name_value; |
| 1514 | SetAttrValue( |
| 1515 | input_info[gather_input_resource_slot].src_node->name(), &feature_name_value); |
| 1516 | |
| 1517 | AttrValue feature_name_to_id_value; |
| 1518 | Status s_feature_to_id = GetFeature2IdAttr( |
| 1519 | input_info[gather_input_resource_slot].src_node->name(), |
| 1520 | &feature_name_to_id_value); |
| 1521 | if (!s_feature_to_id.ok()) return s_feature_to_id; |
| 1522 | |
| 1523 | // get resource shape attr |
| 1524 | int dim_len_value = 0; |
| 1525 | Status s_get_dim = GetShapeValue( |
| 1526 | input_info[gather_input_resource_slot].src_node, &dim_len_value); |
| 1527 | if (!s_get_dim.ok()) return s_get_dim; |
| 1528 | |
| 1529 | AttrValue dim_len_value_int; |
| 1530 | SetAttrValue(dim_len_value, &dim_len_value_int); |
| 1531 |
nothing calls this directly
no test coverage detected