Merge pull request #49383 from ashahab/abin-load-segfault-r2.5

Resolves coredump caused by `tf.data.experimental.save` with prefetch

tensorflow/core/kernels/data/experimental/io_ops.cc

@@ -253,7 +253,11 @@ class LoadDatasetOp::Dataset : public DatasetBase {
     explicit Iterator(const Params& params)
         : DatasetIterator<Dataset>(params) {}
 
-    ~Iterator() override { input_->Unref(); }
+    ~Iterator() override {
+      if (input_) {
+        input_->Unref();
+      }
+    }
 
     Status Initialize(IteratorContext* ctx) override {
       mutex_lock l(mu_);
@@ -331,7 +335,7 @@ class LoadDatasetOp::Dataset : public DatasetBase {
     }
 
     mutex mu_;
-    DatasetBase* input_ TF_GUARDED_BY(mu_);
+    DatasetBase* input_ TF_GUARDED_BY(mu_) = nullptr;
     std::unique_ptr<IteratorBase> input_impl_ TF_GUARDED_BY(mu_);
     std::unique_ptr<InstantiatedCapturedFunction> instantiated_captured_func_;
   };

tensorflow/python/data/experimental/kernel_tests/io_test.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
 import os
 import shutil
 
@@ -111,6 +112,20 @@ def testOptionalElementSpec(self):
     dataset_loaded = io.load(self._test_dir)
     self.assertDatasetsEqual(dataset, dataset_loaded)
 
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRepeatAndPrefetch(self):
+    """This test reproduces github.com/tensorflow/tensorflow/issues/49165"""
+    dataset1 = dataset_ops.Dataset.from_tensor_slices(np.random.rand(16, 32))
+    io.save(dataset1, self._test_dir)
+    dataset = io.load(self._test_dir)
+    dataset = dataset.shuffle(buffer_size=16)
+    dataset = dataset.batch(16)
+    dataset = dataset.repeat()
+    dataset = dataset.prefetch(1)
+    next_element = self.getNext(dataset)
+    for _ in range(30):
+      self.evaluate(next_element())
+
 
 if __name__ == "__main__":
   test.main()