| 1937 | } |
| 1938 | |
| 1939 | Status SavedModelOptimizer::ConvertToHashTableOps() { |
| 1940 | |
| 1941 | for (Node* node : graph_.nodes()) { |
| 1942 | // Get input edges |
| 1943 | std::vector<SrcInfo> input_info; |
| 1944 | Status s_input_info = GetInputNodesInfo(&input_info, node); |
| 1945 | TF_RETURN_IF_ERROR(s_input_info); |
| 1946 | |
| 1947 | Status s_replace; |
| 1948 | if (node->op_def().name() == "KvVarHandleOp") { |
| 1949 | // KvVarHandleOp -> MutableHashTableOfTensorsV2 |
| 1950 | s_replace = ConvertToHashTableOp(node, input_info); |
| 1951 | } else if (node->op_def().name() == "KvResourceGather") { |
| 1952 | // KvResourceGather -> LookupTableFindV2 |
| 1953 | s_replace = ConvertToHashLookupOp(node, input_info); |
| 1954 | } else if (node->op_def().name() == "KvResourceImportV2") { |
| 1955 | // KvResourceImportV2 -> LookupTableImportV2 |
| 1956 | s_replace = ConvertToHashImportOp(node, input_info); |
| 1957 | } |
| 1958 | |
| 1959 | if (!s_replace.ok()) { |
| 1960 | return tensorflow::errors::Internal( |
| 1961 | "Replace Kv ops with hash-table, hash-lookup" |
| 1962 | "or hash-import ops failed, ", |
| 1963 | s_replace.error_message()); |
| 1964 | } |
| 1965 | } |
| 1966 | |
| 1967 | return Status::OK(); |
| 1968 | } |
| 1969 | |
| 1970 | namespace { |
| 1971 | |