[tf.data] Optimization checkpointing improvements.

This CL:
- changes the `OptimizeDataset` checkpointing logic to checkpoint the optimized dataset (as opposed to the original dataset + the optimizations, re-running optimization every time a checkpoint is restored)
- replaces `OpKernelContext` with newly introduced `SerializationContext` in the signature of `AsGraphDefInternal` to reduce the scope of the context and also simplify the logic for overriding the `FunctionLibraryDefinition` when optimizations take place

PiperOrigin-RevId: 208282562

1 files changed, 7 insertions, 5 deletions

diff --git a/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
index f245fc402d..bcf0adacc7 100644
--- a/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
+++ b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
@@ -106,12 +106,14 @@ class GroupByReducerDatasetOp : public UnaryDatasetOpKernel {
     }
 
    protected:
-    Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b,
+    Status AsGraphDefInternal(SerializationContext* ctx,
+                              DatasetGraphDefBuilder* b,
                               Node** output) const override {
-      TF_RETURN_IF_ERROR(b->AddFunction(ctx, key_func().name()));
-      TF_RETURN_IF_ERROR(b->AddFunction(ctx, init_func().name()));
-      TF_RETURN_IF_ERROR(b->AddFunction(ctx, reduce_func().name()));
-      TF_RETURN_IF_ERROR(b->AddFunction(ctx, finalize_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx->flib_def(), key_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx->flib_def(), init_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx->flib_def(), reduce_func().name()));
+      TF_RETURN_IF_ERROR(
+          b->AddFunction(ctx->flib_def(), finalize_func().name()));
       Node* input_graph_node = nullptr;
       TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node));