From 68ad29a9953f7a163c644112737e35571f7ba473 Mon Sep 17 00:00:00 2001
From: Rachel Lim <rachelim@google.com>
Date: Fri, 5 Jun 2020 15:14:39 -0700
Subject: [PATCH] Handle DeadlineExceeded in OptimizeDatasetOp::MakeDataset
 correctly. Before this change, if we encounter a DeadlineExceeded error, the
 output dataset will not be set (i.e. be a nullptr), resulting in "Read
 uninitialized Dataset variant" when we try to use it.

PiperOrigin-RevId: 315005357
Change-Id: I91f1f3b4e99ddc733ec11b5040c6d01b094422ba
---
 tensorflow/core/kernels/data/optimize_dataset_op.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/core/kernels/data/optimize_dataset_op.cc b/tensorflow/core/kernels/data/optimize_dataset_op.cc
index d5d4e8c8f14..c976a8f7b08 100644
--- a/tensorflow/core/kernels/data/optimize_dataset_op.cc
+++ b/tensorflow/core/kernels/data/optimize_dataset_op.cc
@@ -60,6 +60,10 @@ void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
   if (errors::IsDeadlineExceeded(s)) {
     // Ignore DeadlineExceeded as it implies that the attempted rewrite took too
     // long which should not prevent further computation.
+    LOG(WARNING) << s.ToString();
+
+    *output = input;
+    input->Ref();
     return;
   }
   OP_REQUIRES_OK(ctx, s);