| 412 | } |
| 413 | |
| 414 | Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input, |
| 415 | std::function<RewriterConfig(void)> config_factory, |
| 416 | bool optimize_function_library, |
| 417 | DatasetBase** rewritten_input) { |
| 418 | SerializationContext::Params params; |
| 419 | std::vector<std::pair<string, Tensor>> input_list; |
| 420 | params.input_list = &input_list; |
| 421 | params.check_external_state = false; |
| 422 | params.fail_if_unimplemented = false; |
| 423 | params.serialize_data_tensors = false; |
| 424 | SerializationContext serialization_ctx(params); |
| 425 | GraphDef graph_def; |
| 426 | TF_RETURN_IF_ERROR( |
| 427 | AsGraphDef(ctx, input, std::move(serialization_ctx), &graph_def)); |
| 428 | |
| 429 | string output_node; |
| 430 | for (const auto& node : graph_def.node()) { |
| 431 | if (node.op() == "_Retval") { |
| 432 | output_node = node.input(0); |
| 433 | } |
| 434 | } |
| 435 | |
| 436 | VLOG(3) << "Before graph rewrites: " << graph_def.DebugString(); |
| 437 | TF_RETURN_IF_ERROR(ApplyRewrites(ctx, config_factory, |
| 438 | optimize_function_library, &graph_def, |
| 439 | &output_node)); |
| 440 | VLOG(3) << "After graph rewrites: " << graph_def.DebugString(); |
| 441 | |
| 442 | // Instantiate the optimized input pipeline by running the optimized graph |
| 443 | // using the optimized function library. |
| 444 | FunctionLibraryRuntime* flr = nullptr; |
| 445 | std::unique_ptr<ProcessFunctionLibraryRuntime> pflr = nullptr; |
| 446 | std::unique_ptr<FunctionLibraryDefinition> lib_def = nullptr; |
| 447 | TF_RETURN_IF_ERROR( |
| 448 | ctx->function_library()->Clone(&lib_def, &pflr, &flr, true)); |
| 449 | |
| 450 | // Some functions may have been modified without having their names |
| 451 | // changed (for example, nested dataset graphs from FlatMap or |
| 452 | // Interleave). |
| 453 | TF_RETURN_IF_ERROR(AddToFunctionLibrary(lib_def.get(), graph_def.library())); |
| 454 | |
| 455 | Graph graph(OpRegistry::Global()); |
| 456 | TF_RETURN_IF_ERROR(ImportGraphDef({}, graph_def, &graph, nullptr)); |
| 457 | std::vector<Tensor> outputs; |
| 458 | GraphRunner graph_runner(flr->device()); |
| 459 | |
| 460 | TF_RETURN_IF_ERROR( |
| 461 | graph_runner.Run(&graph, flr, input_list, {output_node}, &outputs)); |
| 462 | TF_RETURN_IF_ERROR(GetDatasetFromVariantTensor(outputs[0], rewritten_input)); |
| 463 | (*rewritten_input)->Ref(); |
| 464 | return Status::OK(); |
| 465 | } |
| 466 | |
| 467 | Status VerifyTypesMatch(const DataTypeVector& expected, |
| 468 | const DataTypeVector& received) { |
no test coverage detected