| 2058 | } |
| 2059 | |
| 2060 | Status SavedModelOptimizer::CreateIncrRestoreOp( |
| 2061 | Node* import_op, Node** restore_op) { |
| 2062 | NodeDef incr_restore_def; |
| 2063 | incr_restore_def.set_name(import_op->name() + "_IncrRestore"); |
| 2064 | incr_restore_def.set_op("IncrRestore"); |
| 2065 | DataTypeVector dtypes; |
| 2066 | // Tkeys, dtype, int64 (for version) |
| 2067 | AttrValue* key_attr = nullptr; |
| 2068 | Status s_attr = GetNodeAttr(import_op, "Tkeys", &key_attr); |
| 2069 | TF_RETURN_IF_ERROR(s_attr); |
| 2070 | AttrValue* dtype_attr = nullptr; |
| 2071 | s_attr = GetNodeAttr(import_op, "dtype", &dtype_attr); |
| 2072 | TF_RETURN_IF_ERROR(s_attr); |
| 2073 | dtypes.push_back(key_attr->type()); |
| 2074 | dtypes.push_back(dtype_attr->type()); |
| 2075 | dtypes.push_back(tensorflow::DataType::DT_INT64); |
| 2076 | |
| 2077 | AttrValue attr_value; |
| 2078 | SetAttrValue(dtypes, &attr_value); |
| 2079 | (*incr_restore_def.mutable_attr())["dtypes"] = attr_value; |
| 2080 | Status status; |
| 2081 | *restore_op = graph_.AddNode(incr_restore_def, &status); |
| 2082 | TF_RETURN_IF_ERROR(status); |
| 2083 | |
| 2084 | std::vector<SrcInfo> src_info; |
| 2085 | Status s_src_info = GetInputNodesInfo(&src_info, import_op); |
| 2086 | TF_RETURN_IF_ERROR(s_src_info); |
| 2087 | |
| 2088 | // Add input egdes |
| 2089 | std::vector<SrcInfo> input_info; |
| 2090 | // input: prefix |
| 2091 | input_info.push_back({src_info[import_input_prefix_slot].src_node, |
| 2092 | src_info[import_input_prefix_slot].src_slot}); |
| 2093 | |
| 2094 | // input: tensor_names |
| 2095 | // concat key/value/version strings to one tensor |
| 2096 | Node* kvv_concat_dim_op = nullptr; |
| 2097 | Tensor kvv_concat_dim_val(DT_INT32, TensorShape({})); |
| 2098 | kvv_concat_dim_val.scalar<int>()() = 0; |
| 2099 | TF_RETURN_IF_ERROR( |
| 2100 | CreateFakeConstOp(import_op->name() + "_IncrRestore/kvv_concat_dim", |
| 2101 | DataType::DT_INT32, &graph_, &kvv_concat_dim_op, |
| 2102 | &kvv_concat_dim_val)); |
| 2103 | std::vector<SrcInfo> kvv_concat_inputs; |
| 2104 | kvv_concat_inputs.push_back({kvv_concat_dim_op, 0}); |
| 2105 | |
| 2106 | std::string tensor_kind_names[3] = {"keys", "values", "versions"}; |
| 2107 | for (int i = 0; i < 3; ++i) { |
| 2108 | Tensor suffix(DT_STRING, TensorShape({1})); |
| 2109 | suffix.flat<string>()(0) = "-" + tensor_kind_names[i]; |
| 2110 | |
| 2111 | Node* suffix_node = nullptr; |
| 2112 | TF_RETURN_IF_ERROR( |
| 2113 | CreateFakeConstOp(import_op->name() + "_IncrRestore/" + |
| 2114 | tensor_kind_names[i] + "_suffix", |
| 2115 | DataType::DT_STRING, &graph_, &suffix_node, &suffix)); |
| 2116 | |
| 2117 | std::vector<SrcInfo> join_inputs; |
nothing calls this directly
no test coverage detected