| 223 | } |
| 224 | |
| 225 | void Compute(OpKernelContext* context) override { |
| 226 | const Tensor& prefix = context->input(0); |
| 227 | const string& prefix_string = prefix.scalar<string>()(); |
| 228 | const Tensor& tensor_names = context->input(1); |
| 229 | const Tensor& shape_and_slices = context->input(2); |
| 230 | const Tensor& is_sparse_tensor = context->input(3); |
| 231 | const bool& is_sparse = is_sparse_tensor.scalar<bool>()(); |
| 232 | const auto& shape_and_slices_flat = shape_and_slices.flat<string>(); |
| 233 | const int num_tensors = static_cast<int>(tensor_names.NumElements()); |
| 234 | if (is_sparse) { |
| 235 | BundleReader reader(Env::Default(), prefix_string); |
| 236 | OP_REQUIRES_OK(context, reader.status()); |
| 237 | VLOG(1) << "BundleReader incr, prefix_string: " << prefix_string; |
| 238 | LOG(INFO) << "BundleReader incr, prefix_string: " << prefix_string; |
| 239 | const auto& tensor_names_flat = tensor_names.flat<string>(); |
| 240 | if (num_tensors > 1) { |
| 241 | // EV |
| 242 | if (num_tensors != 3) { |
| 243 | OP_REQUIRES_OK(context, errors::InvalidArgument( |
| 244 | "Incr cpkt restore for ev must has 3 tensors, actually ", |
| 245 | num_tensors, " given")); |
| 246 | } |
| 247 | |
| 248 | const string& ev_keys_name = tensor_names_flat(0); |
| 249 | string incr_tensor_name = ev_keys_name.substr(0, |
| 250 | ev_keys_name.find("-keys")); |
| 251 | // 1 read keys, values and versions |
| 252 | TensorShape incr_shape; |
| 253 | Tensor* incr_keys_tensor = nullptr; |
| 254 | Tensor* incr_values_tensor = nullptr; |
| 255 | Tensor* incr_versions_tensor = nullptr; |
| 256 | OP_REQUIRES_OK(context, |
| 257 | reader.LookupTensorShape(incr_tensor_name + "-sparse_incr_keys", |
| 258 | &incr_shape)); |
| 259 | OP_REQUIRES_OK(context, |
| 260 | context->allocate_output(0, incr_shape, &incr_keys_tensor)); |
| 261 | OP_REQUIRES_OK(context, |
| 262 | reader.Lookup(incr_tensor_name + "-sparse_incr_keys", |
| 263 | incr_keys_tensor)); |
| 264 | OP_REQUIRES_OK(context, |
| 265 | reader.LookupTensorShape(incr_tensor_name + "-sparse_incr_values", |
| 266 | &incr_shape)); |
| 267 | OP_REQUIRES_OK(context, |
| 268 | context->allocate_output(1, incr_shape, &incr_values_tensor)); |
| 269 | OP_REQUIRES_OK(context, |
| 270 | reader.Lookup(incr_tensor_name + "-sparse_incr_values", |
| 271 | incr_values_tensor)); |
| 272 | |
| 273 | OP_REQUIRES_OK(context, |
| 274 | reader.LookupTensorShape(incr_tensor_name + |
| 275 | "-sparse_incr_versions", &incr_shape)); |
| 276 | OP_REQUIRES_OK(context, |
| 277 | context->allocate_output(2, incr_shape, |
| 278 | &incr_versions_tensor)); |
| 279 | OP_REQUIRES_OK(context, |
| 280 | reader.Lookup(incr_tensor_name + "-sparse_incr_versions", |
| 281 | incr_versions_tensor)); |
| 282 | } else { |
nothing calls this directly
no test coverage detected