chore: import upstream snapshot with attribution

tensorflow/python/data/kernel_tests/as_numpy_iterator_test.py

@@ -0,0 +1,144 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.numpy()`."""
+import collections
+import os
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import test
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+
+
+class AsNumpyIteratorTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testBasic(self):
+    ds = dataset_ops.Dataset.range(3)
+    self.assertEqual([0, 1, 2], list(ds.as_numpy_iterator()))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testImmutable(self):
+    ds = dataset_ops.Dataset.from_tensors([1, 2, 3])
+    arr = next(ds.as_numpy_iterator())
+    with self.assertRaisesRegex(ValueError,
+                                'assignment destination is read-only'):
+      arr[0] = 0
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testNestedStructure(self):
+    point = collections.namedtuple('Point', ['x', 'y'])
+    ds = dataset_ops.Dataset.from_tensor_slices({
+        'a': ([1, 2], [3, 4]),
+        'b': [5, 6],
+        'c': point([7, 8], [9, 10])
+    })
+    self.assertEqual([{
+        'a': (1, 3),
+        'b': 5,
+        'c': point(7, 9)
+    }, {
+        'a': (2, 4),
+        'b': 6,
+        'c': point(8, 10)
+    }], list(ds.as_numpy_iterator()))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testNonEager(self):
+    ds = dataset_ops.Dataset.range(10)
+    with self.assertRaises(RuntimeError):
+      ds.as_numpy_iterator()
+
+  def _testInvalidElement(self, element):
+    ds = dataset_ops.Dataset.from_tensors(element)
+    with self.assertRaisesRegex(TypeError,
+                                'is not supported for datasets that'):
+      ds.as_numpy_iterator()
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSparseElement(self):
+    st = sparse_tensor.SparseTensor(
+        indices=[(0, 0), (1, 1), (2, 2)], values=[1, 2, 3], dense_shape=(3, 3))
+    ds = dataset_ops.Dataset.from_tensor_slices(st)
+    dt = sparse_ops.sparse_tensor_to_dense(st)
+    self.assertAllEqual(list(ds.as_numpy_iterator()), dt.numpy())
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRaggedElement(self):
+    lst = [[1, 2], [3], [4, 5, 6]]
+    rt = ragged_factory_ops.constant([lst])
+    # This dataset consists of exactly one ragged tensor.
+    ds = dataset_ops.Dataset.from_tensor_slices(rt)
+    expected = np.array([
+        np.array([1, 2], dtype=np.int32),
+        np.array([3], dtype=np.int32),
+        np.array([4, 5, 6], dtype=np.int32)
+    ], dtype=object)
+    for actual in ds.as_numpy_iterator():
+      self.assertEqual(len(actual), len(expected))
+      for actual_arr, expected_arr in zip(actual, expected):
+        self.assertTrue(np.array_equal(actual_arr, expected_arr),
+                        f'{actual_arr} != {expected_arr}')
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDatasetElement(self):
+    self._testInvalidElement(dataset_ops.Dataset.range(3))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testNestedNonTensorElement(self):
+    tuple_elem = (constant_op.constant([1, 2, 3]), dataset_ops.Dataset.range(3))
+    self._testInvalidElement(tuple_elem)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testNoneElement(self):
+    ds = dataset_ops.Dataset.from_tensors((2, None))
+    self.assertAllEqual(list(ds.as_numpy_iterator()), [(2, None)])
+
+  @combinations.generate(combinations.times(
+      test_base.eager_only_combinations(),
+      combinations.combine(enable_async_ckpt=[True, False])
+  ))
+  def testCompatibleWithCheckpoint(self, enable_async_ckpt):
+    ds = dataset_ops.Dataset.range(10)
+    iterator = ds.as_numpy_iterator()
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_enable_async_checkpoint=enable_async_ckpt)
+    for _ in range(5):
+      next(iterator)
+
+    prefix = os.path.join(self.get_temp_dir(), 'ckpt')
+    save_path = ckpt.save(prefix, options=ckpt_options)
+    self.assertEqual(5, next(iterator))
+    self.assertEqual(6, next(iterator))
+    restore_iter = ds.as_numpy_iterator()
+    restore_ckpt = trackable_utils.Checkpoint(iterator=restore_iter)
+    if enable_async_ckpt:
+      ckpt.sync()  # Otherwise save may not finish yet
+    restore_ckpt.restore(save_path)
+    self.assertEqual(5, next(restore_iter))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/data/kernel_tests/batch_test.py

@@ -0,0 +1,606 @@
+# -*- coding: utf-8 -*-
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.batch()`."""
+
+import time
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+class BatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              count=[0, 28],
+              batch_size=[14, 15],
+              drop_remainder=[True, False],
+              num_parallel_calls=[None, 1, 2, 4])))
+  def testBasic(self, count, batch_size, drop_remainder, num_parallel_calls):
+    """Tests the batch dataset logic for various input configurations.
+
+    Args:
+      count: the number of input elements
+      batch_size: the batch size
+      drop_remainder: whether a smaller batch size should be produced if batch
+        size does not divide number of inputs evenly
+      num_parallel_calls: the number batches to process asynchronously in
+        parallel
+    """
+
+    # The pipeline is TensorSliceDataset -> MapDataset(square_3) ->
+    # RepeatDataset(count) -> BatchDataset(batch_size).
+    components = (np.arange(7),
+                  np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
+                  np.array(37.0) * np.arange(7))
+
+    def _map_fn(x, y, z):
+      return math_ops.square(x), math_ops.square(y), math_ops.square(z)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+        _map_fn).repeat(count).batch(batch_size, drop_remainder,
+                                     num_parallel_calls)
+    get_next = self.getNext(dataset)
+
+    if drop_remainder:
+      dim0 = batch_size
+    else:
+      dim0 = None
+    self.assertEqual(
+        [ts.as_list() for ts in nest.flatten(
+            dataset_ops.get_legacy_output_shapes(dataset))],
+        [[dim0] + list(c.shape[1:]) for c in components])
+
+    num_full_batches = (count * 7) // batch_size
+    for i in range(num_full_batches):
+      result = self.evaluate(get_next())
+      for component, result_component in zip(components, result):
+        for j in range(batch_size):
+          self.assertAllEqual(component[(i * batch_size + j) % 7]**2,
+                              result_component[j])
+    if not drop_remainder and (count * 7) % batch_size > 0:
+      result = self.evaluate(get_next())
+      for component, result_component in zip(components, result):
+        for j in range((count * 7) % batch_size):
+          self.assertAllEqual(
+              component[(num_full_batches * batch_size + j) % 7]**2,
+              result_component[j])
+    with self.assertRaises(errors.OutOfRangeError):
+      result = self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidBatchSize(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = (dataset_ops.Dataset.range(10).batch(0))
+      self.evaluate(dataset._variant_tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLargeBatchSize(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    ).batch(dtypes.int64.max, drop_remainder=False)
+    element = self.evaluate(dataset.get_single_element())
+    self.assertAllEqual(element, list(range(10)))
+    self.assertEqual(element.shape, (10,))
+    self.assertEqual(element.dtype, dtypes.int32)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataset(self):
+
+    def map_fn(i):
+      return dataset_ops.Dataset.from_tensors(i)
+
+    dataset = dataset_ops.Dataset.range(10).map(map_fn).batch(5)
+    dataset = dataset.map(lambda x: x)
+    dataset = dataset.unbatch().flat_map(lambda x: x)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+  def testSparse(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).batch(5)
+    expected_output = [
+        sparse_tensor.SparseTensorValue(
+            indices=[[0, 0], [1, 0], [2, 0], [3, 0], [4, 0]],
+            values=[i * 5, i * 5 + 1, i * 5 + 2, i * 5 + 3, i * 5 + 4],
+            dense_shape=[5, 1]) for i in range(2)
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparseWithDifferentDenseShapes(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=array_ops.expand_dims(
+              math_ops.range(i, dtype=dtypes.int64), 1),
+          values=array_ops.fill([math_ops.cast(i, dtypes.int32)], i),
+          dense_shape=[i])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).batch(5)
+    expected_output = []
+    for i in range(2):
+      expected_indices = []
+      expected_outputs = []
+      for j in range(5):
+        for k in range(i * 5 + j):
+          expected_indices.append([j, k])
+          expected_outputs.append(i * 5 + j)
+      expected_output.append(
+          sparse_tensor.SparseTensorValue(
+              indices=expected_indices,
+              values=expected_outputs,
+              dense_shape=[5, (i + 1) * 5 - 1]))
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparseNested(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).batch(5).batch(2)
+    expected_output = [
+        sparse_tensor.SparseTensorValue(
+            indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], [0, 4, 0],
+                     [1, 0, 0], [1, 1, 0], [1, 2, 0], [1, 3, 0], [1, 4, 0]],
+            values=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            dense_shape=[2, 5, 1])
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShapeError(self):
+
+    def generator():
+      yield [1.0, 2.0, 3.0]
+      yield [4.0, 5.0, 6.0]
+      yield [7.0, 8.0, 9.0, 10.0]
+
+    dataset = (
+        dataset_ops.Dataset.from_generator(
+            generator, dtypes.float32, output_shapes=[None]).batch(3))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(
+            errors.InvalidArgumentError,
+            r"Cannot batch tensors with different shapes in component 0. First "
+            r"element had shape \[3\] and element 2 had shape \[4\]."))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRagged(self):
+
+    def _ragged(i):
+      return ragged_tensor.RaggedTensor.from_tensor(i * [[1]])
+
+    dataset = dataset_ops.Dataset.range(10).map(_ragged).batch(5)
+    expected_output = [
+        ragged_factory_ops.constant([[[0]], [[1]], [[2]], [[3]], [[4]]]),
+        ragged_factory_ops.constant([[[5]], [[6]], [[7]], [[8]], [[9]]])
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRaggedWithDifferentShapes(self):
+    dataset = dataset_ops.Dataset.range(10).map(ragged_math_ops.range).batch(5)
+    expected_output = [
+        ragged_concat_ops.stack([ragged_math_ops.range(i) for i in range(5)]),
+        ragged_concat_ops.stack(
+            [ragged_math_ops.range(i) for i in range(5, 10)])
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRaggedNested(self):
+
+    def _ragged(i):
+      return ragged_tensor.RaggedTensor.from_tensor(i * [[1]])
+
+    dataset = dataset_ops.Dataset.range(10).map(_ragged).batch(5).batch(2)
+    expected_output = [
+        ragged_factory_ops.constant([[[[0]], [[1]], [[2]], [[3]], [[4]]],
+                                     [[[5]], [[6]], [[7]], [[8]], [[9]]]])
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoneComponent(self):
+    dataset = dataset_ops.Dataset.range(10).map(lambda x: (x, None)).batch(
+        10).map(lambda x, y: x)
+    self.assertDatasetProduces(dataset, expected_output=[list(range(10))])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              local_determinism=[None, True, False],
+              global_determinism=[True, False])))
+  def testDeterminismConfiguration(self, local_determinism, global_determinism):
+    expect_determinism = local_determinism or (local_determinism is None and
+                                               global_determinism)
+    elements = list(range(100))
+
+    def dataset_fn(delay_ms):
+
+      def sleep(x):
+        time.sleep(delay_ms / 1000)
+        return x
+
+      def map_function(x):
+        if math_ops.equal(x, 0):
+          return script_ops.py_func(sleep, [x], x.dtype)
+        else:
+          return x
+
+      dataset = dataset_ops.Dataset.from_tensor_slices(elements)
+      dataset = dataset.map(
+          map_function, num_parallel_calls=2, deterministic=local_determinism)
+      dataset = dataset.batch(
+          batch_size=6, num_parallel_calls=2,
+          deterministic=local_determinism).unbatch()
+      opts = options_lib.Options()
+      opts.deterministic = global_determinism
+      dataset = dataset.with_options(opts)
+      return dataset
+
+    self.checkDeterminism(dataset_fn, expect_determinism, elements)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCheckpointLargeBatches(self):
+    # Batches of size 512M
+    dataset = dataset_ops.Dataset.from_tensors(
+        array_ops.ones((64, 1024, 1024), dtype=dtypes.float32)).repeat()
+    dataset = dataset.batch(2, num_parallel_calls=5)
+    iterator = iter(dataset)
+    next(iterator)  # request an element to fill the buffer
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, self.get_temp_dir(), max_to_keep=1)
+    manager.save()
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 1])))
+  def testName(self, num_parallel_calls):
+    dataset = dataset_ops.Dataset.range(5).batch(
+        5, num_parallel_calls=num_parallel_calls, name="batch")
+    self.assertDatasetProduces(dataset, [list(range(5))])
+
+
+class BatchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                          parameterized.TestCase):
+
+  def _build_dataset(self,
+                     multiplier=15.0,
+                     tensor_slice_len=2,
+                     batch_size=2,
+                     num_parallel_calls=None,
+                     options=None):
+    components = (np.arange(tensor_slice_len), np.array([[1, 2, 3]]) *
+                  np.arange(tensor_slice_len)[:, np.newaxis],
+                  np.array(multiplier) * np.arange(tensor_slice_len))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.batch(batch_size, num_parallel_calls=num_parallel_calls)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              symbolic_checkpoint=[False, True], num_parallel_calls=[None, 4]
+          ),
+      )
+  )
+  def test(self, verify_fn, symbolic_checkpoint, num_parallel_calls):
+    tensor_slice_len = 8
+    batch_size = 2
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    num_outputs = tensor_slice_len // batch_size
+    verify_fn(
+        self,
+        lambda: self._build_dataset(
+            15.0, tensor_slice_len, batch_size, num_parallel_calls, options
+        ),
+        num_outputs,
+    )
+
+  def _sparse(self, i):
+    return sparse_tensor.SparseTensorValue(
+        indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+  def _build_dataset_sparse(self, batch_size=5):
+    return dataset_ops.Dataset.range(10).map(self._sparse).batch(batch_size)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testSparse(self, verify_fn):
+    verify_fn(self, self._build_dataset_sparse, num_outputs=2)
+
+  def _build_dataset_nested_sparse(self):
+    return dataset_ops.Dataset.range(10).map(self._sparse).batch(5).batch(2)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testNestedSparse(self, verify_fn):
+    verify_fn(self, self._build_dataset_nested_sparse, num_outputs=1)
+
+
+class BatchRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 2, 3, 4])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4]).batch(2)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).batch(2)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 0))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              count=[0, 10, 20, 30, 40, 50],
+              batch_size=[1, 3, 5, 7, 10, 20],
+              drop_remainder=[True, False])))
+  def testBasic(self, count, batch_size, drop_remainder):
+    """Tests the batch dataset logic for various input configurations.
+
+    Args:
+      count: the number of input elements
+      batch_size: the batch size
+      drop_remainder: whether a smaller batch size should be produced if batch
+        size does not divide number of inputs evenly
+    """
+    dataset = dataset_ops.Dataset.from_tensor_slices(list(range(count))).batch(
+        batch_size=batch_size, drop_remainder=drop_remainder)
+    num_full_batches = count // batch_size
+    for i in range(num_full_batches):
+      expected_batch = np.arange(
+          i * batch_size, (i * batch_size + batch_size), 1, dtype=np.int32)
+      self.assertAllEqual(expected_batch,
+                          self.evaluate(random_access.at(dataset, i)))
+    has_remainder = (not drop_remainder) and (count % batch_size != 0)
+    if has_remainder:
+      expected_batch = np.arange(batch_size * num_full_batches, count, 1)
+      self.assertAllEqual(
+          expected_batch,
+          self.evaluate(random_access.at(dataset, num_full_batches)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(
+          random_access.at(
+              dataset, index=num_full_batches + (1 if has_remainder else 0)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRandomAccessBatchWithShuffle(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6, 7])
+    shuffle_dataset = dataset.shuffle(buffer_size=10, seed=2)
+    batch_dataset = shuffle_dataset.batch(2)
+
+    expected_output = [
+        np.array([5, 2], dtype=np.int32),
+        np.array([4, 7], dtype=np.int32),
+        np.array([1, 3], dtype=np.int32),
+        np.array([6], dtype=np.int32)
+    ]
+    for i in range(4):
+      self.assertAllEqual(expected_output[i],
+                          self.evaluate(random_access.at(batch_dataset, i)))
+
+    # Checks the order is consistent with shuffle dataset.
+    for i in range(3):
+      self.assertAllEqual(
+          expected_output[i][0],
+          self.evaluate(random_access.at(shuffle_dataset, i * 2)))
+      self.assertAllEqual(
+          expected_output[i][1],
+          self.evaluate(random_access.at(shuffle_dataset, (i * 2) + 1)))
+
+    # Checks the remainder is the last element in shuffled dataset.
+    self.assertAllEqual(expected_output[3][0],
+                        self.evaluate(random_access.at(shuffle_dataset, 6)))
+
+
+class BatchGlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              batch_size=[2, 7])))
+  def testBatch(
+      self, dataset_range: int, batch_size: int):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.batch(batch_size, drop_remainder=True)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = global_shuffle_op._global_shuffle(dataset)
+    dataset = dataset.unbatch()
+
+    expected = list(range(0, (dataset_range // batch_size) * batch_size))
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              batch_size=[2, 7],
+              reshuffle=[True, False],
+              seed=[None, 42])))
+  def testReshuffleRepeatEpochs(
+      self,
+      dataset_range: int,
+      batch_size: int,
+      reshuffle: bool,
+      seed: Optional[int]):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.batch(batch_size, drop_remainder=True)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle)
+    dataset = dataset.repeat(2)
+    dataset = dataset.unbatch()
+
+    expected = list(range(0, (dataset_range // batch_size) * batch_size))
+    len_per_iteration = len(expected)
+    expected *= 2
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(output, expected)
+    output_per_iteration = [
+        output[i : i + len_per_iteration]
+        for i in range(0, len(output), len_per_iteration)]
+    if reshuffle:
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[1])
+    else:
+      self.assertEqual(output_per_iteration[0], output_per_iteration[1])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              batch_size=[2, 7])))
+  def testNoDropRemainder(
+      self, dataset_range: int, batch_size: int):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.batch(batch_size, drop_remainder=False)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        "does not support global shuffling with `drop_remainder=False`."):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class BatchGlobalShuffleCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                       parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              batch_size=[2, 3],
+              symbolic_checkpoint=[True, False])))
+  def testBatch(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      batch_size: int,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.batch(batch_size, drop_remainder=True)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(dataset, seed=42)
+      dataset = dataset.unbatch()
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=(dataset_range // batch_size) * batch_size,
+        assert_items_equal=True)
+
+  # Creating multiple iterators with the same seed is only supported in v2 API.
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(tf_api_version=2, mode="eager"),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              batch_size=[2, 3],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testReshuffleEachIteration(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      batch_size: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.batch(batch_size, drop_remainder=True)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      dataset = dataset.unbatch()
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=(dataset_range // batch_size) * batch_size,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/bucket_by_sequence_length_test.py

@@ -0,0 +1,496 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.bucket_by_sequence_length()."""
+import random
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+def _element_length_fn(x, y=None):
+  del y
+  return array_ops.shape(x)[0]
+
+
+def _to_sparse_tensor(record):
+  return sparse_tensor.SparseTensor(**record)
+
+
+def _format_record(array, sparse):
+  if sparse:
+    return {
+        "values": array,
+        "indices": [[i] for i in range(len(array))],
+        "dense_shape": (len(array),)
+    }
+  return array
+
+
+def _get_record_type(sparse):
+  if sparse:
+    return {
+        "values": dtypes.int64,
+        "indices": dtypes.int64,
+        "dense_shape": dtypes.int64
+    }
+  return dtypes.int32
+
+
+def _get_record_shape(sparse):
+  if sparse:
+    return {
+        "values": tensor_shape.TensorShape([
+            None,
+        ]),
+        "indices": tensor_shape.TensorShape([None, 1]),
+        "dense_shape": tensor_shape.TensorShape([
+            1,
+        ])
+    }
+  return tensor_shape.TensorShape([None])
+
+
+class BucketBySequenceLengthTest(test_base.DatasetTestBase,
+                                 parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(param_no_padding=[True, False])))
+  def testBucketDropReminder(self, param_no_padding):
+
+    boundaries = [10, 20, 30]
+    batch_sizes = [10, 8, 4, 2]
+    lengths = [8, 13, 25, 35]
+
+    n_bucket_elements = [28, 7, 6, 5]
+    n_expected_batches = 5
+
+    # Expected sequence lengths of the individual batches.
+    expected_lengths = []
+
+    # Expected sum of all batches with an equal sequence length.
+    # <seq-length>: <expected-total-sum>
+    expected_sums = {}
+
+    # Expected batch sizes of batches depending on the sequence length.
+    # <seq-length>: [batch1_size, ..., batchN_size]
+    expected_batch_sizes = {}
+
+    for length, batch_size, bucket_elements in zip(lengths, batch_sizes,
+                                                   n_bucket_elements):
+      # Calculate the expected sum across all batches of a specific sequence
+      # length.
+      expected_sums[length] = \
+          (bucket_elements - bucket_elements % batch_size) * length
+      # Calculate the expected occurrence of individual batch sizes.
+      expected_batch_sizes[length] = \
+          [batch_size] * (bucket_elements // batch_size)
+      # Calculate the expected occurrence of individual sequence lengths.
+      expected_lengths.extend([length] * (bucket_elements // batch_size))
+
+    def build_dataset(sparse):
+
+      def _generator():
+        # Produce 1 batch for each bucket
+        elements = []
+        for bucket_elements, length in zip(n_bucket_elements, lengths):
+          # Using only full sequences (opposed to the strategy employed in
+          # `testBucket`) makes checking the sum a lot easier.
+          record_len = length
+          for _ in range(bucket_elements):
+            elements.append([1] * record_len)
+        random.shuffle(elements)
+        for el in elements:
+          yield (_format_record(el, sparse),)
+
+      dataset = dataset_ops.Dataset.from_generator(_generator,
+                                                   (_get_record_type(sparse),),
+                                                   (_get_record_shape(sparse),))
+      if sparse:
+        dataset = dataset.map(lambda x: (_to_sparse_tensor(x),))
+      return dataset
+
+    def _test_bucket_by_padding(no_padding):
+      dataset = build_dataset(sparse=no_padding)
+      dataset = dataset.bucket_by_sequence_length(
+          element_length_func=_element_length_fn,
+          bucket_boundaries=boundaries,
+          bucket_batch_sizes=batch_sizes,
+          no_padding=no_padding,
+          drop_remainder=True)
+
+      get_next = self.getNext(dataset)
+      batches = []
+      for _ in range(n_expected_batches):
+        batch, = self.evaluate(get_next())
+        batches.append(batch)
+
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+      generated_lengths = []
+
+      # <seq-length>: <total-sum>
+      generated_sums = {}
+
+      # <seq-length>: [<batch_size>, ...]
+      generated_batch_sizes = {}
+
+      for length, batch_size, bucket_elements in zip(lengths, batch_sizes,
+                                                     n_bucket_elements):
+        # Initialize the sum across all batches.
+        generated_sums[length] = 0
+        # Initialize the individual batch sizes.
+        generated_batch_sizes[length] = []
+
+      for batch in batches:
+        shape = batch.dense_shape if no_padding else batch.shape
+        length = shape[1]
+        generated_lengths.append(length)
+
+        batch_size = shape[0]
+        generated_batch_sizes[length].append(batch_size)
+
+        batch_sum = batch.values.sum() if no_padding else batch.sum()
+        generated_sums[length] += batch_sum
+
+      for l in lengths:
+        # Make sure the sum of the batch contents is correct for the individual
+        # sequence lengths.
+        self.assertEqual(
+            generated_sums[l], expected_sums[l], "Tensor sums did not match! "
+            "expected: {}, generated: {}".format(expected_sums, generated_sums))
+
+        # Make sure the individual batch sizes are generated as expected.
+        self.assertEqual(
+            sorted(generated_batch_sizes[l]), sorted(expected_batch_sizes[l]),
+            "Batch-sizes did not match! "
+            "expected: {}, generated: {}".format(
+                sorted(expected_batch_sizes[l]),
+                sorted(generated_batch_sizes[l])))
+
+      # Make sure the generated sequence lengths appear as often as expected.
+      self.assertEqual(
+          sorted(generated_lengths), sorted(expected_lengths),
+          "The generated sequence lengths did not match! "
+          "expected: {}, generated: {}".format(
+              sorted(expected_lengths), sorted(generated_lengths)))
+
+    _test_bucket_by_padding(param_no_padding)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(param_no_padding=[True, False])))
+  def testBucket(self, param_no_padding):
+
+    boundaries = [10, 20, 30]
+    batch_sizes = [10, 8, 4, 2]
+    lengths = [8, 13, 25, 35]
+
+    def build_dataset(sparse):
+
+      def _generator():
+        # Produce 1 batch for each bucket
+        elements = []
+        for batch_size, length in zip(batch_sizes, lengths):
+          record_len = length - 1
+          for _ in range(batch_size):
+            elements.append([1] * record_len)
+            record_len = length
+        random.shuffle(elements)
+        for el in elements:
+          yield (_format_record(el, sparse),)
+
+      dataset = dataset_ops.Dataset.from_generator(_generator,
+                                                   (_get_record_type(sparse),),
+                                                   (_get_record_shape(sparse),))
+      if sparse:
+        dataset = dataset.map(lambda x: (_to_sparse_tensor(x),))
+      return dataset
+
+    def _test_bucket_by_padding(no_padding):
+      dataset = build_dataset(sparse=no_padding)
+      dataset = dataset.bucket_by_sequence_length(
+          element_length_func=_element_length_fn,
+          bucket_boundaries=boundaries,
+          bucket_batch_sizes=batch_sizes,
+          no_padding=no_padding)
+      get_next = self.getNext(dataset)
+      batches = []
+      for _ in range(4):
+        batch, = self.evaluate(get_next())
+        batches.append(batch)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+      batch_sizes_val = []
+      lengths_val = []
+      for batch in batches:
+        shape = batch.dense_shape if no_padding else batch.shape
+        batch_size = shape[0]
+        length = shape[1]
+        batch_sizes_val.append(batch_size)
+        lengths_val.append(length)
+        if not context.executing_eagerly():
+          sum_check = batch.values.sum() if no_padding else batch.sum()
+          self.assertEqual(sum_check, batch_size * length - 1)
+      self.assertEqual(sum(batch_sizes_val), sum(batch_sizes))
+      self.assertEqual(sorted(batch_sizes), sorted(batch_sizes_val))
+      self.assertEqual(sorted(lengths), sorted(lengths_val))
+
+    _test_bucket_by_padding(param_no_padding)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPadToBoundary(self):
+
+    boundaries = [10, 20, 30]
+    batch_sizes = [10, 8, 4, 2]
+    lengths = [8, 13, 25]
+
+    def element_gen():
+      # Produce 1 batch for each bucket
+      elements = []
+      for batch_size, length in zip(batch_sizes[:-1], lengths):
+        for _ in range(batch_size):
+          elements.append([1] * length)
+      random.shuffle(elements)
+      for el in elements:
+        yield (el,)
+      for _ in range(batch_sizes[-1]):
+        el = [1] * (boundaries[-1] + 5)
+        yield (el,)
+
+    element_len = lambda el: array_ops.shape(el)[0]
+    dataset = dataset_ops.Dataset.from_generator(element_gen, (dtypes.int64,),
+                                                 ([None],))
+    dataset = dataset.bucket_by_sequence_length(
+        element_length_func=element_len,
+        bucket_boundaries=boundaries,
+        bucket_batch_sizes=batch_sizes,
+        pad_to_bucket_boundary=True)
+    get_next = self.getNext(dataset)
+
+    batches = []
+    for _ in range(3):
+      batch, = self.evaluate(get_next())
+      batches.append(batch)
+    with self.assertRaisesOpError("bucket_boundaries"):
+      self.evaluate(get_next())
+
+    batch_sizes_val = []
+    lengths_val = []
+    for batch in batches:
+      batch_size = batch.shape[0]
+      length = batch.shape[1]
+      batch_sizes_val.append(batch_size)
+      lengths_val.append(length)
+    batch_sizes = batch_sizes[:-1]
+    self.assertEqual(sum(batch_sizes_val), sum(batch_sizes))
+    self.assertEqual(sorted(batch_sizes), sorted(batch_sizes_val))
+    self.assertEqual([boundary - 1 for boundary in sorted(boundaries)],
+                     sorted(lengths_val))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPadToBoundaryNoExtraneousPadding(self):
+
+    boundaries = [3, 7, 11]
+    batch_sizes = [2, 2, 2, 2]
+    lengths = range(1, 11)
+
+    def element_gen():
+      for length in lengths:
+        yield ([1] * length,)
+
+    element_len = lambda element: array_ops.shape(element)[0]
+    dataset = dataset_ops.Dataset.from_generator(element_gen, (dtypes.int64,),
+                                                 ([None],))
+    dataset = dataset.bucket_by_sequence_length(
+        element_length_func=element_len,
+        bucket_boundaries=boundaries,
+        bucket_batch_sizes=batch_sizes,
+        pad_to_bucket_boundary=True)
+    get_next = self.getNext(dataset)
+
+    batches = []
+    for _ in range(5):
+      batch, = self.evaluate(get_next())
+      batches.append(batch)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    self.assertAllEqual(batches[0], [[1, 0], [1, 1]])
+    self.assertAllEqual(batches[1], [[1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 0, 0]])
+    self.assertAllEqual(batches[2], [[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1]])
+    self.assertAllEqual(
+        batches[3],
+        [[1, 1, 1, 1, 1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
+    self.assertAllEqual(
+        batches[4],
+        [[1, 1, 1, 1, 1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(param_no_padding=[True, False])))
+  def testTupleElements(self, param_no_padding):
+
+    def build_dataset(sparse):
+
+      def _generator():
+        text = [[1, 2, 3], [3, 4, 5, 6, 7], [1, 2], [8, 9, 0, 2, 3]]
+        label = [1, 2, 1, 2]
+        for x, y in zip(text, label):
+          yield (_format_record(x, sparse), y)
+
+      dataset = dataset_ops.Dataset.from_generator(
+          generator=_generator,
+          output_types=(_get_record_type(sparse), dtypes.int32),
+          output_shapes=(_get_record_shape(sparse),
+                         tensor_shape.TensorShape([])))
+      if sparse:
+        dataset = dataset.map(lambda x, y: (_to_sparse_tensor(x), y))
+      return dataset
+
+    def _test_tuple_elements_by_padding(no_padding):
+      dataset = build_dataset(sparse=no_padding)
+      dataset = dataset.bucket_by_sequence_length(
+          element_length_func=_element_length_fn,
+          bucket_batch_sizes=[2, 2, 2],
+          bucket_boundaries=[0, 8],
+          no_padding=no_padding)
+      shapes = dataset_ops.get_legacy_output_shapes(dataset)
+      self.assertEqual([None, None], shapes[0].as_list())
+      self.assertEqual([None], shapes[1].as_list())
+
+    _test_tuple_elements_by_padding(param_no_padding)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(param_drop_remainder=[True, False])))
+  def testBucketSparse(self, param_drop_remainder):  # pylint: disable=g-doc-args
+    """Tests bucketing of sparse tensors (case where `no_padding` == True).
+
+    Test runs on following dataset:
+      [
+        [0],
+        [0, 1],
+        [0, 1, 2]
+        ...
+        [0, ..., max_len - 1]
+      ]
+    Sequences are bucketed by length and batched with
+      `batch_size` < `bucket_size`.
+    """
+
+    min_len = 0
+    max_len = 100
+    batch_size = 7
+    bucket_size = 10
+
+    def _build_dataset():
+      input_data = [range(i + 1) for i in range(min_len, max_len)]
+
+      def generator_fn():
+        for record in input_data:
+          yield _format_record(record, sparse=True)
+
+      dataset = dataset_ops.Dataset.from_generator(
+          generator=generator_fn, output_types=_get_record_type(sparse=True))
+      dataset = dataset.map(_to_sparse_tensor)
+      return dataset
+
+    def _compute_expected_batches(drop_remainder):
+      """Computes expected batch outputs and stores in a set."""
+      all_expected_sparse_tensors = set()
+      for bucket_start_len in range(min_len, max_len, bucket_size):
+        if drop_remainder:
+          batch_offsets = [0]
+        else:
+          batch_offsets = range(0, bucket_size, batch_size)
+
+        for batch_offset in batch_offsets:
+          batch_start_len = bucket_start_len + batch_offset
+          batch_end_len = min(batch_start_len + batch_size,
+                              bucket_start_len + bucket_size)
+          expected_indices = []
+          expected_values = []
+          for length in range(batch_start_len, batch_end_len):
+            for val in range(length + 1):
+              expected_indices.append((length - batch_start_len, val))
+              expected_values.append(val)
+          expected_sprs_tensor = (tuple(expected_indices),
+                                  tuple(expected_values))
+          all_expected_sparse_tensors.add(expected_sprs_tensor)
+      return all_expected_sparse_tensors
+
+    def _compute_batches(dataset):
+      """Computes actual batch outputs of dataset and stores in a set."""
+      batch = self.getNext(dataset)
+      all_sparse_tensors = set()
+      with self.assertRaises(errors.OutOfRangeError):
+        while True:
+          output = self.evaluate(batch())
+          sprs_tensor = (tuple([tuple(idx) for idx in output.indices]),
+                         tuple(output.values))
+          all_sparse_tensors.add(sprs_tensor)
+
+      return all_sparse_tensors
+
+    dataset = _build_dataset()
+    boundaries = range(min_len + bucket_size + 1, max_len, bucket_size)
+    dataset = dataset.bucket_by_sequence_length(
+        element_length_func=_element_length_fn,
+        bucket_boundaries=boundaries,
+        bucket_batch_sizes=[batch_size] * (len(boundaries) + 1),
+        no_padding=True,
+        drop_remainder=param_drop_remainder)
+    batches = _compute_batches(dataset)
+    expected_batches = _compute_expected_batches(param_drop_remainder)
+    self.assertEqual(batches, expected_batches)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCardinality(self):
+
+    boundaries = [3, 7, 11]
+    batch_sizes = [2, 2, 2, 2]
+    lengths = range(1, 11)
+
+    def element_gen():
+      for length in lengths:
+        yield ([1] * length,)
+
+    element_len = lambda element: array_ops.shape(element)[0]
+    dataset = dataset_ops.Dataset.from_generator(element_gen, (dtypes.int64,),
+                                                 ([None],)).repeat()
+    dataset = dataset.bucket_by_sequence_length(
+        element_length_func=element_len,
+        bucket_boundaries=boundaries,
+        bucket_batch_sizes=batch_sizes,
+        pad_to_bucket_boundary=True)
+    self.assertEqual(self.evaluate(dataset.cardinality()), dataset_ops.INFINITE)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/cache_test.py

@@ -0,0 +1,765 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.cache()`."""
+import functools
+import os
+from os import path
+import shutil
+import tempfile
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class FileCacheTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(FileCacheTest, self).setUp()
+    self.tmp_dir = tempfile.mkdtemp()
+    self.cache_prefix = path.join(self.tmp_dir, "cache")
+
+  def tearDown(self):
+    if self.tmp_dir:
+      shutil.rmtree(self.tmp_dir, ignore_errors=True)
+    super(FileCacheTest, self).tearDown()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheDatasetPassthrough(self):
+    components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]),
+                  np.array([9.0, 10.0, 11.0, 12.0]))
+
+    def dataset_fn(count=5, filename=None):
+      repeat_dataset = (
+          dataset_ops.Dataset.from_tensor_slices(components).repeat(count))
+      if filename:
+        return repeat_dataset.cache(filename)
+      else:
+        return repeat_dataset
+
+    self.assertEqual(
+        tuple([c.shape[1:] for c in components]),
+        dataset_ops.get_legacy_output_shapes(dataset_fn()))
+
+    get_next = self.getNext(dataset_fn())
+
+    # First run without caching to collect the "ground truth".
+    elements = []
+    for _ in range(20):
+      elements.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Assert that the cached dataset has the same elements as the
+    # "ground truth".
+    get_next = self.getNext(dataset_fn(filename=self.cache_prefix))
+    cached_elements = []
+    for _ in range(20):
+      cached_elements.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertAllEqual(elements, cached_elements)
+
+    # Re-initialize with an empty upstream (to throw errors.OutOfRangeError
+    # if we didn't use the cache).
+    get_next = self.getNext(dataset_fn(count=0, filename=self.cache_prefix))
+    replayed_elements = []
+    for _ in range(20):
+      replayed_elements.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertEqual(cached_elements, replayed_elements)
+
+    # Re-initialize with an empty upstream and a missing cache file (should
+    # throw errors.OutOfRangeError immediately).
+    get_next = self.getNext(
+        dataset_fn(count=0, filename=self.cache_prefix + "nonsense"))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcurrentWriters(self):
+    components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]),
+                  np.array([9.0, 10.0, 11.0, 12.0]))
+
+    cache_dataset1 = (
+        dataset_ops.Dataset.from_tensor_slices(components).cache(
+            self.cache_prefix))
+    cache_dataset2 = (
+        dataset_ops.Dataset.from_tensor_slices(components).cache(
+            self.cache_prefix))
+
+    get_next1 = self.getNext(cache_dataset1)
+    get_next2 = self.getNext(cache_dataset2)
+
+    self.evaluate(get_next1())  # this should succeed
+
+    with self.assertRaises(errors.AlreadyExistsError):
+      self.evaluate(get_next2())
+
+    self.evaluate(get_next1())  # this should continue to succeed
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcurrentReaders(self):
+    components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]),
+                  np.array([9.0, 10.0, 11.0, 12.0]))
+
+    cache_dataset1 = (
+        dataset_ops.Dataset.from_tensor_slices(components).cache(
+            self.cache_prefix))
+    cache_dataset2 = (
+        dataset_ops.Dataset.from_tensor_slices(components).cache(
+            self.cache_prefix))
+
+    get_next1 = self.getNext(cache_dataset1)
+    get_next2 = self.getNext(cache_dataset2)
+
+    elements = []
+    for _ in range(4):
+      elements.append(self.evaluate(get_next1()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next1())
+
+    # Re-initialize
+    get_next1 = self.getNext(cache_dataset1, requires_initialization=True)
+    get_next2 = self.getNext(cache_dataset2, requires_initialization=True)
+
+    # Reading concurrently should succeed.
+    elements_itr1 = []
+    elements_itr2 = []
+    elements_itr2.append(self.evaluate(get_next2()))
+    elements_itr1.append(self.evaluate(get_next1()))
+    elements_itr2.append(self.evaluate(get_next2()))
+    elements_itr1.append(self.evaluate(get_next1()))
+    # Intentionally reversing the order
+    elements_itr1.append(self.evaluate(get_next1()))
+    elements_itr2.append(self.evaluate(get_next2()))
+    elements_itr1.append(self.evaluate(get_next1()))
+    elements_itr2.append(self.evaluate(get_next2()))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next2())
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next1())
+
+    self.assertAllEqual(elements, elements_itr1)
+    self.assertAllEqual(elements, elements_itr2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadingPastEndOfSequence(self):
+    dataset = dataset_ops.Dataset.range(10).cache(self.cache_prefix)
+    dataset = dataset.map(lambda a: a).batch(4).repeat(2)
+    expected_output = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]] * 2
+    self.assertDatasetProduces(dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheZipped(self):
+    def make_dataset(i):
+      cache_path = self.cache_prefix + "_" + str(i)
+      return dataset_ops.Dataset.range(100).shuffle(100).cache(cache_path)
+
+    datasets = [make_dataset(i) for i in range(3)]
+    dataset = dataset_ops.Dataset.zip(tuple(datasets))
+    first_order = self.getDatasetOutput(dataset)
+    second_order = self.getDatasetOutput(dataset)
+    self.assertEqual(first_order, second_order)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCleaningUpCacheFiles(self):
+
+    def do_test(i):
+      dataset = dataset_ops.Dataset.range(10).cache(self.cache_prefix)
+      get_next = self.getNext(dataset)
+      for _ in range(i):
+        try:
+          self.evaluate(get_next())
+        except errors.OutOfRangeError:
+          break
+
+    if not context.executing_eagerly():
+      self.skipTest(
+          "Test requires eager mode for iterators to be deconstructed")
+
+    for i in [0, 3, 10, 12, 15]:
+      do_test(i)
+
+
+class MemoryCacheTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheDatasetPassthrough(self):
+    with ops.device("cpu:0"):
+      repeat_count = variables.Variable(constant_op.constant(10, dtypes.int64))
+      dataset = dataset_ops.Dataset.range(3).flat_map(
+          lambda x: dataset_ops.Dataset.from_tensors(x).repeat(repeat_count))
+
+      options = options_lib.Options()
+      options.experimental_optimization.inject_prefetch = False
+      dataset = dataset.with_options(options)
+
+      cached_dataset = dataset.cache().repeat(2)
+      uncached_dataset = dataset.repeat(2)
+
+      self.evaluate(repeat_count.initializer)
+      # Needs to be initializable to capture the variable.
+      cached_next = self.getNext(cached_dataset, requires_initialization=True)
+      uncached_next = self.getNext(
+          uncached_dataset, requires_initialization=True)
+      for i in range(3):
+        for _ in range(10):
+          self.assertEqual(self.evaluate(cached_next()), i)
+          self.assertEqual(self.evaluate(uncached_next()), i)
+
+      self.evaluate(repeat_count.assign(0))
+
+      # The uncached iterator should now be empty.
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(uncached_next())
+
+      # The cached iterator replays from cache.
+      for i in range(3):
+        for _ in range(10):
+          self.assertEqual(self.evaluate(cached_next()), i)
+
+      # The cached iterator should now be empty.
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(cached_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyCacheReading(self):
+    components = (np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]),
+                  np.array([9.0, 10.0, 11.0, 12.0]))
+
+    repeat_dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).repeat(0))
+    cache_dataset = repeat_dataset.cache()
+
+    self.assertDatasetProduces(cache_dataset, expected_output=[])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcurrentReaders(self):
+
+    dataset_fn = lambda: dataset_ops.Dataset.range(5).cache()
+    d1 = dataset_fn().map(lambda x: x + 1)
+    d2 = dataset_fn().map(lambda x: x + 6)
+
+    get_next1 = self.getNext(d1)
+
+    self.assertEqual(1, self.evaluate(get_next1()))
+    self.assertEqual(2, self.evaluate(get_next1()))
+    self.assertEqual(3, self.evaluate(get_next1()))
+
+    get_next2 = self.getNext(d2)
+
+    self.assertEqual(6, self.evaluate(get_next2()))
+    self.assertEqual(7, self.evaluate(get_next2()))
+    self.assertEqual(4, self.evaluate(get_next1()))  # interleave execution
+    self.assertEqual([8, 5],
+                     [self.evaluate(get_next2()),
+                      self.evaluate(get_next1())])
+    self.assertEqual(9, self.evaluate(get_next2()))
+    self.assertEqual(10, self.evaluate(get_next2()))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next2())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next1())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheTakeRepeat(self):
+    dataset = dataset_ops.Dataset.range(10).cache().take(5).repeat(2)
+
+    expected_output = [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheRepeatEpochs(self):
+    counter = variables.Variable(0)
+    self.evaluate(counter.initializer)
+
+    def increment_fn(x):
+      counter.assign_add(1)
+      return x
+
+    dataset = dataset_ops.Dataset.range(10).map(increment_fn).cache().repeat(2)
+    options = options_lib.Options()
+    options.experimental_optimization.inject_prefetch = False
+    dataset = dataset.with_options(options)
+    get_next = self.getNext(dataset, requires_initialization=True)
+
+    # first epoch
+    for i in range(10):
+      self.assertEqual(i, self.evaluate(counter))
+      self.assertEqual(i, self.evaluate(get_next()))
+    # second epoch
+    for i in range(10):
+      self.assertEqual(10, self.evaluate(counter))
+      self.assertEqual(i, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testCacheIterationEpochs(self):
+    counter = variables.Variable(0)
+    self.evaluate(counter.initializer)
+
+    def increment_fn(x):
+      counter.assign_add(1)
+      return x
+
+    dataset = dataset_ops.Dataset.range(10).map(increment_fn).cache()
+    options = options_lib.Options()
+    options.experimental_optimization.inject_prefetch = False
+    dataset = dataset.with_options(options)
+
+    # first epoch
+    i = 0
+    for elem in dataset:
+      self.assertEqual(i, self.evaluate(elem))
+      i += 1
+      self.assertEqual(i, self.evaluate(counter))
+
+    # second epoch
+    i = 0
+    for elem in dataset:
+      self.assertEqual(10, self.evaluate(counter))
+      self.assertEqual(i, self.evaluate(elem))
+      i += 1
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testCacheV2ResourceCapture(self):
+
+    def make_dataset():
+      ids = dataset_ops.Dataset.range(10)
+      ids = ids.cache()
+
+      def interleave_fn(dataset, _):
+        return dataset
+
+      dataset = dataset_ops.Dataset.range(1)
+      dataset = dataset.interleave(functools.partial(interleave_fn, ids))
+      return dataset
+
+    results = []
+    for elem in make_dataset():
+      results.append(elem.numpy())
+
+    self.assertAllEqual(results, range(10))
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testCacheV2ConcurrentIterators(self):
+
+    dataset = dataset_ops.Dataset.range(10).cache()
+
+    it1 = iter(dataset)
+    it2 = iter(dataset)
+
+    for i in range(10):
+      self.assertEqual(next(it1), i)
+      self.assertEqual(next(it2), i)
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testCacheKnownCardinality(self):
+
+    # Check that a dataset which produces random permutation of range(10) ends
+    # up being cached when we read all of its element but do not reach EOF.
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.shuffle(10, reshuffle_each_iteration=True).cache()
+
+    it = iter(dataset)
+
+    results = []
+    for _ in range(10):
+      results.append(next(it))
+
+    it = iter(dataset)
+    for i in range(10):
+      self.assertEqual(next(it), results[i])
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCheckpointFinishedCache(self):
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = ds.cache()
+
+    iterator = iter(ds)
+    for i in range(num_elements):
+      self.assertEqual(next(iterator).numpy(), i)
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, self.get_temp_dir(), max_to_keep=1)
+    manager.save()
+    manager.restore_or_initialize()
+    with self.assertRaises(StopIteration):
+      next(iterator)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).cache(name="cache")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class CacheCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                          parameterized.TestCase):
+
+  def setUp(self):
+    super(CacheCheckpointTest, self).setUp()
+    self.range_size = 10
+    self.num_repeats = 3
+    self.num_outputs = self.range_size * self.num_repeats
+    self.cache_file_prefix = "test"
+
+  def make_dataset_fn(self, is_memory):
+    if is_memory:
+      filename = ""
+    else:
+      filename = os.path.join(self.get_temp_dir(), self.cache_file_prefix)
+
+    def ds_fn():
+      return dataset_ops.Dataset.range(self.range_size).cache(filename).repeat(
+          self.num_repeats)
+
+    return ds_fn
+
+  def expected_outputs(self):
+    return list(range(self.range_size)) * self.num_repeats
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointBeforeOneEpoch(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Generate 5 entries from iterator and save checkpoint.
+    outputs = self.gen_outputs(ds_fn, [], 5, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, range(5))
+
+    # Restore from checkpoint and produce the rest of the elements from the
+    # iterator.
+    outputs.extend(
+        self.gen_outputs(
+            ds_fn, [],
+            self.num_outputs - 5,
+            ckpt_saved=True,
+            verify_exhausted=False))
+    self.assertSequenceEqual(outputs, self.expected_outputs())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointBeforeOneEpochThenRunFewSteps(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Generate 8 entries from iterator but save checkpoint after producing 5.
+    outputs = self.gen_outputs(
+        ds_fn, [5], 8, verify_exhausted=False, save_checkpoint_at_end=False)
+    self.assertSequenceEqual(outputs, range(8))
+
+    outputs = outputs[:5]
+    outputs.extend(
+        self.gen_outputs(
+            ds_fn, [],
+            self.num_outputs - 5,
+            ckpt_saved=True,
+            verify_exhausted=False))
+    self.assertSequenceEqual(outputs, self.expected_outputs())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointAfterOneEpoch(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Generate 15 entries from iterator and save checkpoint.
+    outputs = self.gen_outputs(ds_fn, [], 15, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) + list(range(5)))
+
+    # Restore from checkpoint and produce the rest of the elements from the
+    # iterator.
+    outputs.extend(
+        self.gen_outputs(
+            ds_fn, [],
+            self.num_outputs - 15,
+            ckpt_saved=True,
+            verify_exhausted=False))
+    self.assertSequenceEqual(outputs, self.expected_outputs())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointAfterOneEpochThenRunFewSteps(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Generate 18 entries from iterator but save checkpoint after producing 15.
+    outputs = self.gen_outputs(
+        ds_fn, [15], 18, verify_exhausted=False, save_checkpoint_at_end=False)
+    self.assertSequenceEqual(outputs, list(range(10)) + list(range(8)))
+
+    outputs = list(range(10)) + list(range(5)) + self.gen_outputs(
+        ds_fn, [],
+        self.num_outputs - 15,
+        ckpt_saved=True,
+        verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) * 3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointBeforeOneEpochButRunCompleteEpoch(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Generate 13 entries from iterator but save checkpoint after producing 5.
+    outputs = self.gen_outputs(
+        ds_fn, [5], 13, verify_exhausted=False, save_checkpoint_at_end=False)
+    self.assertSequenceEqual(outputs, list(range(10)) + list(range(3)))
+
+    # Since we ran for more than one epoch, the cache was completely written.
+    # The ckpt was saved when the iterator was in cache-write mode. Test that
+    # the iterator falls back to read mode after restoring if the cache has
+    # been completely written.
+
+    outputs = list(range(5)) + self.gen_outputs(
+        ds_fn, [],
+        self.num_outputs - 5,
+        ckpt_saved=True,
+        verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) * 3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointUnusedWriterIterator(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Checkpoint before get_next is called even once.
+    outputs = self.gen_outputs(ds_fn, [], 0, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, [])
+
+    outputs = self.gen_outputs(
+        ds_fn, [], self.num_outputs, ckpt_saved=True, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) * 3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testCheckpointUnusedMidwayWriterIterator(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Produce 5 elements and checkpoint.
+    outputs = self.gen_outputs(ds_fn, [], 5, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, range(5))
+
+    # Restore from checkpoint, then produce no elements and checkpoint.
+    outputs.extend(
+        self.gen_outputs(ds_fn, [], 0, ckpt_saved=True, verify_exhausted=False))
+    self.assertSequenceEqual(outputs, range(5))
+
+    # Restore from checkpoint and produce rest of the elements.
+    outputs.extend(
+        self.gen_outputs(
+            ds_fn, [],
+            self.num_outputs - 5,
+            ckpt_saved=True,
+            verify_exhausted=False))
+    self.assertSequenceEqual(outputs, list(range(10)) * 3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testUnusedCheckpointError(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Produce 5 elements and save ckpt.
+    outputs = self.gen_outputs(ds_fn, [], 5, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, range(5))
+
+    if is_memory:
+      outputs = self.gen_outputs(
+          ds_fn, [], self.num_outputs, verify_exhausted=False)
+      self.assertSequenceEqual(outputs, self.expected_outputs())
+    else:
+      # Since the complete cache has not been written, a new iterator which does
+      # not restore the checkpoint will throw an error since there is a partial
+      # cache shard.
+      with self.assertRaises(errors.AlreadyExistsError):
+        outputs = self.gen_outputs(
+            ds_fn, [], self.num_outputs, verify_exhausted=False)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(is_memory=[True, False])))
+  def testIgnoreCheckpointIfCacheWritten(self, is_memory):
+    ds_fn = self.make_dataset_fn(is_memory)
+
+    # Produce 15 elements and save ckpt. This will write the complete cache.
+    outputs = self.gen_outputs(ds_fn, [], 15, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) + list(range(5)))
+
+    # Build the iterator again but do not restore from ckpt. Since the cache
+    # has already been written we should be able to use it.
+    outputs = self.gen_outputs(
+        ds_fn, [], self.num_outputs, verify_exhausted=False)
+    self.assertSequenceEqual(outputs, list(range(10)) * 3)
+
+
+class CacheRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 3, 4])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]).cache()
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheRangeDataset(self):
+    dataset = dataset_ops.Dataset.range(10).cache()
+    expected_elements = list(range(10))
+    self.verifyRandomAccess(dataset, expected_elements)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheOneDimensionalElements(self):
+    tensor = [1, 2, 3]
+    dataset = dataset_ops.Dataset.from_tensor_slices(tensor).cache()
+    self.verifyRandomAccess(dataset, tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheTwoDimensionalElements(self):
+    tensor = [[1, 2], [3, 4]]
+    dataset = dataset_ops.Dataset.from_tensor_slices(tensor).cache()
+    self.verifyRandomAccess(dataset, tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheThreeComponents(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        ([1, 2], [3, 4], [5, 6])).cache()
+    expected = [(1, 3, 5), (2, 4, 6)]
+    self.verifyRandomAccess(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheInputDatasetNotRandomlyAccessible(self):
+    dataset = dataset_ops.Dataset.range(10)
+    initial_state = constant_op.constant(0, dtypes.int64)
+    scan_func = lambda state, i: (state + i, state + i)
+    dataset = dataset.scan(
+        initial_state=initial_state, scan_func=scan_func).cache()
+    expected = [0, 1, 3, 6, 10, 15, 21, 28, 36, 45]
+    self.verifyRandomAccess(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheInputDatasetUnknownCardinality(self):
+    dataset = dataset_ops.Dataset.range(20).filter(
+        lambda x: math_ops.equal(x % 2, 0)).cache()
+    expected = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
+    self.verifyRandomAccess(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheInputDatasetInfiniteCardinality(self):
+    dataset = dataset_ops.Dataset.range(20).filter(
+        lambda x: math_ops.equal(x % 2, 0)).repeat(-1).cache()
+    expected = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 0, 2]
+    # Since the dataset has infinite cardinality, random access with caching
+    # will cache through the requested index. In this case, random access
+    # with caching will cache through index 11.
+    self.verifyRandomAccessInfiniteCardinality(dataset, expected)
+
+
+class CacheGlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              repetitions=[1, 2],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def test(
+      self,
+      dataset_range: int,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.cache()
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(0, dataset_range)) * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(dataset_output, self.evaluate(dataset.cardinality()))
+
+
+class CacheGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              repetitions=[1, 2],
+              reshuffle_each_iteration=[True, False])))
+  def test(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.cache()
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      return global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=dataset_range * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/cardinality_test.py

@@ -0,0 +1,197 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.cardinality()`."""
+import functools
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+# pylint: disable=g-long-lambda
+def _test_combinations():
+  v1_only_cases = [
+      ("Map1", lambda: dataset_ops.Dataset.range(5).map(lambda x: x),
+       dataset_ops.UNKNOWN),
+      ("Map2", lambda: dataset_ops.Dataset.range(5).map(
+          lambda x: x, num_parallel_calls=1), dataset_ops.UNKNOWN),
+  ]
+  v2_only_cases = [
+      ("Map1", lambda: dataset_ops.Dataset.range(5).map(lambda x: x), 5),
+      ("Map2", lambda: dataset_ops.Dataset.range(5).map(
+          lambda x: x, num_parallel_calls=1), 5),
+  ]
+  v1_and_v2_cases = [
+      ("Batch1",
+       lambda: dataset_ops.Dataset.range(5).batch(2, drop_remainder=True), 2),
+      ("Batch2",
+       lambda: dataset_ops.Dataset.range(5).batch(2, drop_remainder=False), 3),
+      ("Batch3",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).batch(2),
+       dataset_ops.UNKNOWN),
+      ("Batch4", lambda: dataset_ops.Dataset.range(5).repeat().batch(2),
+       dataset_ops.INFINITE),
+      ("Cache1", lambda: dataset_ops.Dataset.range(5).cache(), 5),
+      ("Cache2", lambda: dataset_ops.Dataset.range(5).cache("foo"), 5),
+      ("Concatenate1", lambda: dataset_ops.Dataset.range(5).concatenate(
+          dataset_ops.Dataset.range(5)), 10),
+      ("Concatenate2",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).concatenate(
+           dataset_ops.Dataset.range(5)), dataset_ops.UNKNOWN),
+      ("Concatenate3", lambda: dataset_ops.Dataset.range(5).repeat().
+       concatenate(dataset_ops.Dataset.range(5)), dataset_ops.INFINITE),
+      ("Concatenate4", lambda: dataset_ops.Dataset.range(5).concatenate(
+          dataset_ops.Dataset.range(5).filter(lambda _: True)),
+       dataset_ops.UNKNOWN),
+      ("Concatenate5",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).concatenate(
+           dataset_ops.Dataset.range(5).filter(lambda _: True)),
+       dataset_ops.UNKNOWN),
+      ("Concatenate6", lambda: dataset_ops.Dataset.range(5).repeat().
+       concatenate(dataset_ops.Dataset.range(5).filter(lambda _: True)),
+       dataset_ops.INFINITE),
+      ("Concatenate7", lambda: dataset_ops.Dataset.range(5).concatenate(
+          dataset_ops.Dataset.range(5).repeat()), dataset_ops.INFINITE),
+      ("Concatenate8",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).concatenate(
+           dataset_ops.Dataset.range(5).repeat()), dataset_ops.INFINITE),
+      ("Concatenate9",
+       lambda: dataset_ops.Dataset.range(5).repeat().concatenate(
+           dataset_ops.Dataset.range(5).repeat()), dataset_ops.INFINITE),
+      ("FlatMap", lambda: dataset_ops.Dataset.range(5).flat_map(
+          lambda _: dataset_ops.Dataset.from_tensors(0)), dataset_ops.UNKNOWN),
+      ("Filter", lambda: dataset_ops.Dataset.range(5).filter(lambda _: True),
+       dataset_ops.UNKNOWN),
+      ("FromTensors1", lambda: dataset_ops.Dataset.from_tensors(0), 1),
+      ("FromTensors2", lambda: dataset_ops.Dataset.from_tensors((0, 1)), 1),
+      ("FromTensorSlices1",
+       lambda: dataset_ops.Dataset.from_tensor_slices([0, 0, 0]), 3),
+      ("FromTensorSlices2", lambda: dataset_ops.Dataset.from_tensor_slices(
+          ([0, 0, 0], [1, 1, 1])), 3),
+      ("Interleave1", lambda: dataset_ops.Dataset.range(5).interleave(
+          lambda _: dataset_ops.Dataset.from_tensors(0), cycle_length=1),
+       dataset_ops.UNKNOWN),
+      ("Interleave2", lambda: dataset_ops.Dataset.range(5).interleave(
+          lambda _: dataset_ops.Dataset.from_tensors(0),
+          cycle_length=1,
+          num_parallel_calls=1), dataset_ops.UNKNOWN),
+      ("Interleave3", lambda: dataset_ops.Dataset.range(5).repeat().interleave(
+          lambda _: dataset_ops.Dataset.from_tensors(0),
+          cycle_length=1,
+          num_parallel_calls=1), dataset_ops.INFINITE),
+      ("PaddedBatch1", lambda: dataset_ops.Dataset.range(5).padded_batch(
+          2, [], drop_remainder=True), 2),
+      ("PaddedBatch2", lambda: dataset_ops.Dataset.range(5).padded_batch(
+          2, [], drop_remainder=False), 3),
+      ("PaddedBatch3", lambda: dataset_ops.Dataset.range(5).filter(
+          lambda _: True).padded_batch(2, []), dataset_ops.UNKNOWN),
+      ("PaddedBatch4",
+       lambda: dataset_ops.Dataset.range(5).repeat().padded_batch(2, []),
+       dataset_ops.INFINITE),
+      ("Prefetch", lambda: dataset_ops.Dataset.range(5).prefetch(buffer_size=1),
+       5),
+      ("Range1", lambda: dataset_ops.Dataset.range(0), 0),
+      ("Range2", lambda: dataset_ops.Dataset.range(5), 5),
+      ("Range3", lambda: dataset_ops.Dataset.range(5, 10), 5),
+      ("Range4", lambda: dataset_ops.Dataset.range(10, 5), 0),
+      ("Range5", lambda: dataset_ops.Dataset.range(5, 10, 2), 3),
+      ("Range6", lambda: dataset_ops.Dataset.range(10, 5, -2), 3),
+      ("Range7", lambda: dataset_ops.Dataset.range(0, 0, -2), 0),
+      ("Range8", lambda: dataset_ops.Dataset.range(3, 3, 1), 0),
+      ("Range9", lambda: dataset_ops.Dataset.range(-4, -4, 2), 0),
+      ("Range10", lambda: dataset_ops.Dataset.range(1, 0, 3), 0),
+      ("Range11", lambda: dataset_ops.Dataset.range(0, 1, -3), 0),
+      ("Repeat1", lambda: dataset_ops.Dataset.range(0).repeat(0), 0),
+      ("Repeat2", lambda: dataset_ops.Dataset.range(1).repeat(0), 0),
+      ("Repeat3", lambda: dataset_ops.Dataset.range(0).repeat(5), 0),
+      ("Repeat4", lambda: dataset_ops.Dataset.range(1).repeat(5), 5),
+      ("Repeat5", lambda: dataset_ops.Dataset.range(0).repeat(), 0),
+      ("Repeat6", lambda: dataset_ops.Dataset.range(1).repeat(),
+       dataset_ops.INFINITE),
+      ("Shuffle", lambda: dataset_ops.Dataset.range(5).shuffle(buffer_size=1),
+       5),
+      ("Shard1", lambda: dataset_ops.Dataset.range(5).shard(2, 0), 3),
+      ("Shard2", lambda: dataset_ops.Dataset.range(5).shard(8, 7), 0),
+      ("Shard3",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).shard(2, 0),
+       dataset_ops.UNKNOWN),
+      ("Shard4", lambda: dataset_ops.Dataset.range(5).repeat().shard(2, 0),
+       dataset_ops.INFINITE),
+      ("Skip1", lambda: dataset_ops.Dataset.range(5).skip(2), 3),
+      ("Skip2", lambda: dataset_ops.Dataset.range(5).skip(8), 0),
+      ("Skip3",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).skip(2),
+       dataset_ops.UNKNOWN),
+      ("Skip4", lambda: dataset_ops.Dataset.range(5).repeat().skip(2),
+       dataset_ops.INFINITE),
+      ("Take1", lambda: dataset_ops.Dataset.range(5).take(2), 2),
+      ("Take2", lambda: dataset_ops.Dataset.range(5).take(8), 5),
+      ("Take3",
+       lambda: dataset_ops.Dataset.range(5).filter(lambda _: True).take(2),
+       dataset_ops.UNKNOWN),
+      ("Take4", lambda: dataset_ops.Dataset.range(5).repeat().take(2), 2),
+      ("Window1", lambda: dataset_ops.Dataset.range(5).window(
+          size=2, shift=2, drop_remainder=True), 2),
+      ("Window2", lambda: dataset_ops.Dataset.range(5).window(
+          size=2, shift=2, drop_remainder=False), 3),
+      ("Zip1", lambda: dataset_ops.Dataset.zip(dataset_ops.Dataset.range(5)),
+       5),
+      ("Zip2", lambda: dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.range(5), dataset_ops.Dataset.range(3))), 3),
+      ("Zip3", lambda: dataset_ops.Dataset.zip((dataset_ops.Dataset.range(
+          5), dataset_ops.Dataset.range(3).repeat())), 5),
+      ("Zip4", lambda: dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.range(5).repeat(), dataset_ops.Dataset.range(3).
+           repeat())), dataset_ops.INFINITE),
+      ("Zip5", lambda: dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.range(5), dataset_ops.Dataset.range(3).filter(
+              lambda _: True))), dataset_ops.UNKNOWN),
+  ]
+  def reduce_cases_to_combinations(x, y):
+    name, dataset_fn, expected_result = y
+    return x + combinations.combine(
+        dataset_fn=combinations.NamedObject(name, dataset_fn),
+        expected_result=expected_result)
+
+  def cases_to_combinations(cases):
+    return functools.reduce(reduce_cases_to_combinations, cases, [])
+
+  v1_only_combinations = combinations.times(
+      combinations.combine(tf_api_version=1, mode=["eager", "graph"]),
+      cases_to_combinations(v1_only_cases))
+  v2_only_combinations = combinations.times(
+      combinations.combine(tf_api_version=2, mode=["eager", "graph"]),
+      cases_to_combinations(v2_only_cases))
+  v1_and_v2_combinations = combinations.times(
+      combinations.combine(tf_api_version=[1, 2], mode=["eager", "graph"]),
+      cases_to_combinations(v1_and_v2_cases))
+
+  return v1_only_combinations + v2_only_combinations + v1_and_v2_combinations
+
+
+class CardinalityTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for `tf.data.Dataset.cardinality()`."""
+
+  @combinations.generate(_test_combinations())
+  def testCardinality(self, dataset_fn, expected_result):
+    dataset = dataset_fn()
+    self.assertEqual(self.evaluate(dataset.cardinality()), expected_result)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/checkpoint_test.py

@@ -0,0 +1,604 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for checkpointing tf.data iterators."""
+import os
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.data.experimental.ops import grouping
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import scan_ops
+from tensorflow.python.data.experimental.ops import take_while_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import test
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+
+
+# TODO(jsimsa): Add missing test combinations.
+class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def tearDown(self):
+    prefix = self._iterator_checkpoint_prefix()
+    pattern = prefix + "*"
+    files = gfile.Glob(pattern)
+    map(gfile.Remove, files)
+    super(CheckpointTest, self).tearDown()
+
+  def _iterator_checkpoint_prefix(self):
+    return os.path.join(self.get_temp_dir(), "iterator")
+
+  def _save_op(self, iterator_resource):
+    iterator_state_variant = gen_dataset_ops.serialize_iterator(
+        iterator_resource)
+    save_op = io_ops.write_file(
+        self._iterator_checkpoint_prefix(),
+        parsing_ops.serialize_tensor(iterator_state_variant))
+    return save_op
+
+  def _restore_op(self, iterator_resource):
+    iterator_state_variant = parsing_ops.parse_tensor(
+        io_ops.read_file(self._iterator_checkpoint_prefix()), dtypes.variant)
+    restore_op = gen_dataset_ops.deserialize_iterator(iterator_resource,
+                                                      iterator_state_variant)
+    return restore_op
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSaveRestore(self):
+
+    def _build_graph(start, stop):
+      iterator = dataset_ops.make_initializable_iterator(
+          dataset_ops.Dataset.range(start, stop))
+      init_op = iterator.initializer
+      get_next = iterator.get_next()
+      save_op = self._save_op(iterator._iterator_resource)
+      restore_op = self._restore_op(iterator._iterator_resource)
+      return init_op, get_next, save_op, restore_op
+
+    # Saving and restoring in different sessions.
+    start = 2
+    stop = 10
+    break_point = 5
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, _ = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        for i in range(start, break_point):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, _, restore_op = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_point, stop):
+          self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+    # Saving and restoring in same session.
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, restore_op = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        for i in range(start, break_point):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_point, stop):
+          self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testInitThenRestore(self):
+    # Note: Calling init_op before restore_op is redundant. This test just makes
+    # sure we do not fail if restore is called on an already initialized
+    # iterator resource.
+
+    def _build_graph(start, stop):
+      dataset = dataset_ops.Dataset.range(start, stop)
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+      init_op = iterator.initializer
+      get_next = iterator.get_next()
+      save_op = self._save_op(iterator._iterator_resource)
+      restore_op = self._restore_op(iterator._iterator_resource)
+      return init_op, get_next, save_op, restore_op
+
+    # Saving and restoring in different sessions.
+    start = 2
+    stop = 10
+    break_point = 5
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, _ = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        for i in range(start, break_point):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, _, restore_op = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_point, stop):
+          self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testMultipleSaves(self):
+
+    def _build_graph(start, stop):
+      iterator = dataset_ops.make_initializable_iterator(
+          dataset_ops.Dataset.range(start, stop))
+      init_op = iterator.initializer
+      get_next = iterator.get_next()
+      save_op = self._save_op(iterator._iterator_resource)
+      restore_op = self._restore_op(iterator._iterator_resource)
+      return init_op, get_next, save_op, restore_op
+
+    start = 2
+    stop = 10
+    break_point1 = 5
+    break_point2 = 7
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, _ = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        for i in range(start, break_point1):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, restore_op = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_point1, break_point2):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+
+    break_point2 = 7
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, restore_op = _build_graph(start, stop)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_point2, stop):
+          self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSaveRestoreWithRepeat(self):
+
+    def _build_graph(start, stop, num_epochs):
+      iterator = dataset_ops.make_initializable_iterator(
+          dataset_ops.Dataset.range(start, stop).repeat(num_epochs))
+      init_op = iterator.initializer
+      get_next = iterator.get_next()
+      save_op = self._save_op(iterator._iterator_resource)
+      restore_op = self._restore_op(iterator._iterator_resource)
+      return init_op, get_next, save_op, restore_op
+
+    start = 2
+    stop = 10
+    num_epochs = 5
+    break_range = 5
+    break_epoch = 3
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, restore_op = _build_graph(
+          start, stop, num_epochs)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        # Note: There is no checkpoint saved currently so a NotFoundError is
+        # raised.
+        with self.assertRaises(errors.NotFoundError):
+          sess.run(init_op)
+          sess.run(restore_op)
+        for _ in range(break_epoch - 1):
+          for i in range(start, stop):
+            self.assertEqual(i, sess.run(get_next))
+        for i in range(start, break_range):
+          self.assertEqual(i, sess.run(get_next))
+        sess.run(save_op)
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, _, restore_op = _build_graph(start, stop, num_epochs)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        for i in range(break_range, stop):
+          self.assertEqual(i, sess.run(get_next))
+        for _ in range(break_epoch, num_epochs):
+          for i in range(start, stop):
+            self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSaveRestoreExhaustedIterator(self):
+
+    def _build_graph(start, stop, num_epochs):
+      iterator = dataset_ops.make_initializable_iterator(
+          dataset_ops.Dataset.range(start, stop).repeat(num_epochs))
+      init_op = iterator.initializer
+      get_next = iterator.get_next()
+      save_op = self._save_op(iterator._iterator_resource)
+      restore_op = self._restore_op(iterator._iterator_resource)
+      return init_op, get_next, save_op, restore_op
+
+    start = 2
+    stop = 10
+    num_epochs = 5
+    with ops.Graph().as_default() as g:
+      init_op, get_next, save_op, restore_op = _build_graph(
+          start, stop, num_epochs)
+      with self.session(graph=g) as sess:
+        sess.run(variables.global_variables_initializer())
+        sess.run(init_op)
+        # Note: There is no checkpoint saved currently so a NotFoundError is
+        # raised.
+        with self.assertRaises(errors.NotFoundError):
+          sess.run(init_op)
+          sess.run(restore_op)
+        for _ in range(num_epochs):
+          for i in range(start, stop):
+            self.assertEqual(i, sess.run(get_next))
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+        sess.run(save_op)
+
+    with ops.Graph().as_default() as g:
+      init_op, get_next, _, restore_op = _build_graph(start, stop, num_epochs)
+      with self.session(graph=g) as sess:
+        sess.run(init_op)
+        sess.run(restore_op)
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  @combinations.generate(combinations.times(
+      test_base.eager_only_combinations(),
+      combinations.combine(enable_async_ckpt=[True, False])
+  ))
+  def testSaveRestoreOneShotIterator(self, enable_async_ckpt):
+    checkpoint_directory = self.get_temp_dir()
+    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6]).map(
+        math_ops.square).batch(2)
+    iterator = iter(dataset)
+    get_next = iterator.get_next
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_enable_async_checkpoint=enable_async_ckpt)
+    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
+    self.assertAllEqual([1, 4], get_next())
+    save_path = checkpoint.save(checkpoint_prefix, options=ckpt_options)
+    self.assertAllEqual([9, 16], get_next())
+    self.assertAllEqual([25, 36], get_next())
+    checkpoint.restore(save_path).run_restore_ops()
+    self.assertAllEqual([9, 16], get_next())
+    self.assertAllEqual([25, 36], get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      get_next()
+
+  @combinations.generate(combinations.times(
+      test_base.eager_only_combinations(),
+      combinations.combine(enable_async_ckpt=[True, False])
+  ))
+  def testSaveRestoreMultipleIterator(self, enable_async_ckpt):
+    checkpoint_directory = self.get_temp_dir()
+    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+    dataset = dataset.map(math_ops.square).batch(2)
+    iterator_1 = iter(dataset)
+    get_next_1 = iterator_1.get_next
+    iterator_2 = iter(dataset)
+    get_next_2 = iterator_2.get_next
+    dataset_2 = dataset_ops.Dataset.range(10)
+    iterator_3 = iter(dataset_2)
+    get_next_3 = iterator_3.get_next
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_enable_async_checkpoint=enable_async_ckpt)
+    checkpoint = trackable_utils.Checkpoint(
+        iterator_1=iterator_1, iterator_2=iterator_2, iterator_3=iterator_3)
+    self.assertAllEqual([1, 4], get_next_1())
+    self.assertAllEqual(0, get_next_3())
+    self.assertAllEqual(1, get_next_3())
+    self.assertAllEqual(2, get_next_3())
+    save_path = checkpoint.save(checkpoint_prefix, options=ckpt_options)
+    self.assertAllEqual([1, 4], get_next_2())
+    self.assertAllEqual([9, 16], get_next_2())
+    self.assertAllEqual(3, get_next_3())
+    checkpoint.restore(save_path).run_restore_ops()
+    self.assertAllEqual([9, 16], get_next_1())
+    self.assertAllEqual([1, 4], get_next_2())
+    self.assertAllEqual(3, get_next_3())
+
+  @combinations.generate(combinations.times(
+      test_base.eager_only_combinations(),
+      combinations.combine(enable_async_ckpt=[True, False])
+  ))
+  def testRestoreExhaustedIterator(self, enable_async_ckpt):
+    checkpoint_directory = self.get_temp_dir()
+    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+    dataset = dataset_ops.Dataset.range(3)
+    iterator = iter(dataset)
+    get_next = iterator.get_next
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_enable_async_checkpoint=enable_async_ckpt)
+    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
+    self.assertAllEqual(0, get_next())
+    self.assertAllEqual(1, get_next())
+    save_path = checkpoint.save(checkpoint_prefix, options=ckpt_options)
+    self.assertAllEqual(2, get_next())
+    checkpoint.restore(save_path).run_restore_ops()
+    self.assertAllEqual(2, get_next())
+    save_path = checkpoint.save(checkpoint_prefix, options=ckpt_options)
+    checkpoint.restore(save_path).run_restore_ops()
+    with self.assertRaises(errors.OutOfRangeError):
+      get_next()
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRestoreInReconstructedIteratorInitializable(self):
+    checkpoint_directory = self.get_temp_dir()
+    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+    dataset = dataset_ops.Dataset.range(10)
+    iterator = iter(dataset)
+    get_next = iterator.get_next
+    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
+    for i in range(5):
+      checkpoint.restore(
+          checkpoint_management.latest_checkpoint(
+              checkpoint_directory)).initialize_or_restore()
+      for j in range(2):
+        self.assertEqual(i * 2 + j, self.evaluate(get_next()))
+      checkpoint.save(file_prefix=checkpoint_prefix)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveRestoreReshuffleDataset(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.shuffle(10, reshuffle_each_iteration=True)
+    iterator = iter(dataset)
+    ckpt = trackable_utils.Checkpoint(
+        step=variables.Variable(0), iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, self.get_temp_dir(), max_to_keep=3)
+
+    iter1 = [next(iterator).numpy() for _ in range(5)]
+
+    manager.save()
+    iter2 = [next(iterator).numpy() for _ in range(5)]
+
+    ckpt.restore(manager.latest_checkpoint)
+    iter3 = [next(iterator).numpy() for _ in range(5)]
+
+    self.assertNotEqual(iter1, iter2)
+    self.assertCountEqual(iter2, iter3)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveRestoreModifiedDataset(self):
+    ckpt_dir = self.get_temp_dir()
+    dataset = dataset_ops.Dataset.range(10)
+    iterator = iter(dataset)
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, ckpt_dir, max_to_keep=3)
+
+    for _ in range(5):
+      next(iterator)
+    manager.save()
+
+    # Define a different dataset and try to restore into its iterator.
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
+    iterator = iter(dataset)
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, ckpt_dir, max_to_keep=3)
+    with self.assertRaisesRegex(
+        errors.NotFoundError,
+        "Make sure the dataset definition has not changed"):
+      ckpt.restore(manager.latest_checkpoint)
+
+  def _assertNotCheckpointable(self, dataset):
+    iterator = iter(dataset)
+    ckpt = trackable_utils.Checkpoint(
+        step=variables.Variable(0), iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, self.get_temp_dir(), max_to_keep=3)
+    with self.assertRaises(errors.FailedPreconditionError):
+      manager.save()
+
+  @staticmethod
+  def _statefulInt64Func(_):
+    return random_ops.random_uniform((), 0, 1, dtypes.int64)
+
+  @staticmethod
+  def _statefulBoolFunc(_):
+    return random_ops.random_uniform((), 0, 1, dtypes.int64) < 1
+
+  @staticmethod
+  def _statefulDatasetFunc(_):
+    x = random_ops.random_uniform((), 0, 1, dtypes.int64)
+    return dataset_ops.Dataset.range(x)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulFilterNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.filter(self._statefulBoolFunc)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulFlatMapNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.flat_map(self._statefulDatasetFunc)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulInterleaveNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.interleave(self._statefulDatasetFunc)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulMapNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.map(self._statefulBoolFunc)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulParallelInterleaveNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.interleave(
+        self._statefulDatasetFunc, num_parallel_calls=2)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulParallelMapNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.map(self._statefulBoolFunc, num_parallel_calls=2)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulGroupByReducerNotCheckpointable(self):
+    stateful_key_func = self._statefulInt64Func
+    key_func = lambda _: math_ops.cast(0, dtypes.int64)
+    stateful_init_func = self._statefulBoolFunc
+    init_func = lambda x: True
+    stateful_reduce_func = lambda _, x: self._statefulBoolFunc(x)
+    reduce_func = lambda _, x: True
+    stateful_finalize_func = self._statefulBoolFunc
+    finalize_func = lambda x: True
+
+    test_cases = [
+        (stateful_key_func, init_func, reduce_func, finalize_func),
+        (key_func, stateful_init_func, reduce_func, finalize_func),
+        (key_func, init_func, stateful_reduce_func, finalize_func),
+        (key_func, init_func, reduce_func, stateful_finalize_func),
+    ]
+    for key_func, init_func, reduce_func, finalize_func in test_cases:
+      dataset = dataset_ops.Dataset.range(10)
+      reducer = grouping.Reducer(init_func, reduce_func, finalize_func)
+      dataset = dataset.apply(grouping.group_by_reducer(key_func, reducer))
+      self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulGroupByWindowNotCheckpointable(self):
+    stateful_key_func = self._statefulInt64Func
+    key_func = lambda _: math_ops.cast(0, dtypes.int64)
+    stateful_reduce_func = lambda _, x: self._statefulDatasetFunc(x)
+    reduce_func = lambda _, x: x
+    stateful_window_func = self._statefulInt64Func
+    window_func = lambda x: math_ops.cast(0, dtypes.int64)
+
+    test_cases = [
+        (stateful_key_func, reduce_func, window_func),
+        (key_func, stateful_reduce_func, window_func),
+        (key_func, reduce_func, stateful_window_func),
+    ]
+    for key_func_fn, reduce_func_fn, window_func in test_cases:
+      dataset = dataset_ops.Dataset.range(10)
+      dataset = dataset.apply(
+          grouping.group_by_window(
+              key_func_fn, reduce_func_fn, window_size_func=window_func))
+      self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulMapAndBatchNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.map(self._statefulBoolFunc)
+    dataset = dataset.batch(2)
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulExperimentalParallelInterleaveNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(self._statefulDatasetFunc, 2))
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulScanNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+
+    def stateful_scan(state, element):
+      return state, self._statefulBoolFunc(element)
+
+    dataset = dataset.apply(scan_ops.scan(0, stateful_scan))
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulTakeWhileNotCheckpointable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(take_while_ops.take_while(self._statefulBoolFunc))
+    self._assertNotCheckpointable(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStatefulExternalPolicy(self):
+    checkpoint_directory = self.get_temp_dir()
+    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+    dataset = dataset_ops.Dataset.range(4)
+
+    def fn(x):
+      return x * x
+
+    dataset = dataset.map(
+        lambda x: script_ops.eager_py_func(fn, [x], dtypes.int64))
+
+    options = options_lib.Options()
+    options.experimental_external_state_policy = (
+        options_lib.ExternalStatePolicy.WARN)
+    dataset = dataset.with_options(options)
+
+    iterator = iter(dataset)
+    get_next = iterator.get_next
+    checkpoint = trackable_utils.Checkpoint(iterator=iterator)
+    self.assertEqual(0, get_next().numpy())
+    self.assertEqual(1, get_next().numpy())
+    save_path = checkpoint.save(checkpoint_prefix)
+    self.assertEqual(4, get_next().numpy())
+    self.assertEqual(9, get_next().numpy())
+    checkpoint.restore(save_path).run_restore_ops()
+    self.assertEqual(4, get_next().numpy())
+    self.assertEqual(9, get_next().numpy())
+    with self.assertRaises(errors.OutOfRangeError):
+      get_next()
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/checkpoint_test_base.py

@@ -0,0 +1,675 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Base test class for checkpointing datasets."""
+
+import os
+
+import numpy as np
+
+from tensorflow.python.checkpoint import checkpoint as tracking_util
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.util import nest
+
+
+def remove_variants(get_next_op):
+  # TODO(b/72408568): Remove this once session.run can get variant tensors.
+  """Remove variants from a nest structure, so sess.run will execute."""
+
+  def _remove_variant(x):
+    if isinstance(x, tensor.Tensor) and x.dtype == dtypes.variant:
+      return ()
+    else:
+      return x
+
+  return nest.map_structure(_remove_variant, get_next_op)
+
+
+def default_test_combinations():
+  """Returns the default test combinations for testing checkpointing."""
+
+  def disable_optimizations(ds_fn):
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+
+    def ds_fn_no_opt():
+      return ds_fn().with_options(options)
+
+    return ds_fn_no_opt
+
+  def verify_unused_iterator(
+      obj, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    obj.verify_unused_iterator(
+        ds_fn=disable_optimizations(ds_fn=ds_fn),
+        num_outputs=num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  verify_unused_iterator_combination = combinations.combine(
+      verify_fn=combinations.NamedObject(
+          "verify_unused_iterator", verify_unused_iterator))
+
+  def verify_fully_used_iterator(
+      obj, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    obj.verify_fully_used_iterator(
+        ds_fn=disable_optimizations(ds_fn=ds_fn),
+        num_outputs=num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  verify_fully_used_iterator_combination = combinations.combine(
+      verify_fn=combinations.NamedObject(
+          "verify_fully_used_iterator", verify_fully_used_iterator))
+
+  def verify_exhausted_iterator(
+      obj, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    obj.verify_exhausted_iterator(
+        ds_fn=disable_optimizations(ds_fn=ds_fn),
+        num_outputs=num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  verify_exhausted_iterator_combination = combinations.combine(
+      verify_fn=combinations.NamedObject(
+          "verify_exhausted_iterator", verify_exhausted_iterator))
+
+  def verify_multiple_breaks(
+      obj, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    obj.verify_multiple_breaks(
+        ds_fn=disable_optimizations(ds_fn=ds_fn),
+        num_outputs=num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  verify_multiple_breaks_combination = combinations.combine(
+      verify_fn=combinations.NamedObject(
+          "verify_multiple_breaks", verify_multiple_breaks))
+
+  def verify_reset_restored_iterator(
+      obj, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    obj.verify_reset_restored_iterator(
+        ds_fn=disable_optimizations(ds_fn=ds_fn),
+        num_outputs=num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  verify_reset_restored_iterator_combination = combinations.combine(
+      verify_fn=combinations.NamedObject(
+          "verify_reset_restored_iterator", verify_reset_restored_iterator))
+
+  return (verify_unused_iterator_combination +
+          verify_fully_used_iterator_combination +
+          verify_exhausted_iterator_combination +
+          verify_multiple_breaks_combination +
+          verify_reset_restored_iterator_combination)
+
+
+# TODO(b/72657739): Remove sparse_tensor argument, which is to test the
+# (deprecated) saveable `SparseTensorSliceDataset`, once the API
+# `from_sparse_tensor_slices()` and related tests are deleted.
+class CheckpointTestBase(test.TestCase):
+  """Base test class for checkpointing datasets."""
+
+  def tearDown(self):
+    self._delete_ckpt()
+    super(CheckpointTestBase, self).tearDown()
+
+  def verify_unused_iterator(self,
+                             ds_fn,
+                             num_outputs,
+                             sparse_tensors=False,
+                             verify_exhausted=True,
+                             assert_items_equal=False):
+    """Verifies that saving and restoring an unused iterator works.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      verify_exhausted: Whether to verify that the iterator has been exhausted
+        after producing `num_outputs` elements.
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    self.verify_run_with_breaks(
+        ds_fn, [0],
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=verify_exhausted,
+        assert_items_equal=assert_items_equal)
+
+  def verify_fully_used_iterator(self,
+                                 ds_fn,
+                                 num_outputs,
+                                 sparse_tensors=False,
+                                 assert_items_equal=False):
+    """Verifies that saving and restoring a fully used iterator works.
+
+    Note that this only checks saving and restoring an iterator from which
+    `num_outputs` items have been produced but does not check for an
+    exhausted iterator, i.e., one from which an OutOfRange error has been
+    returned.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if test fails.
+    """
+    self.verify_run_with_breaks(
+        ds_fn,
+        [num_outputs],
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        assert_items_equal=assert_items_equal)
+
+  def verify_exhausted_iterator(
+      self, ds_fn, num_outputs, sparse_tensors=False, assert_items_equal=False):
+    """Verifies that saving and restoring an exhausted iterator works.
+
+    An exhausted iterator is one which has returned an OutOfRange error.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    del assert_items_equal
+    self.gen_outputs(
+        ds_fn, [],
+        num_outputs,
+        verify_exhausted=True,
+        sparse_tensors=sparse_tensors)
+    actual = self.gen_outputs(
+        ds_fn, [],
+        0,
+        ckpt_saved=True,
+        verify_exhausted=True,
+        sparse_tensors=sparse_tensors)
+    self.assertLen(actual, 0)
+
+  def verify_multiple_breaks(self,
+                             ds_fn,
+                             num_outputs,
+                             num_breaks=10,
+                             sparse_tensors=False,
+                             verify_exhausted=True,
+                             assert_items_equal=False):
+    """Attempts to save/restore at multiple break points.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      num_breaks: The number of break points. These are uniformly spread in [0,
+        num_outputs] both inclusive.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      verify_exhausted: Whether to verify that the iterator has been exhausted
+        after producing `num_outputs` elements.
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    self.verify_run_with_breaks(
+        ds_fn,
+        self.gen_break_points(num_outputs, num_breaks),
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=verify_exhausted,
+        assert_items_equal=assert_items_equal)
+
+  def verify_reset_restored_iterator(self,
+                                     ds_fn,
+                                     num_outputs,
+                                     break_point=None,
+                                     sparse_tensors=False,
+                                     verify_exhausted=True,
+                                     assert_items_equal=False):
+    """Attempts to re-initialize a restored iterator.
+
+    This is useful when restoring a training checkpoint during validation.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      break_point: Break point. Optional. Defaults to num_outputs/2.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      verify_exhausted: Whether to verify that the iterator has been exhausted
+        after producing `num_outputs` elements.
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    if context.executing_eagerly():
+      self.skipTest("Eager mode iteration do not support re-initialization.")
+
+    break_point = num_outputs // 2 if not break_point else break_point
+
+    # Collect ground truth containing all outputs.
+    expected = self.gen_outputs(
+        ds_fn, [],
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=verify_exhausted)
+
+    # Skip some items and save checkpoint.
+    self.gen_outputs(
+        ds_fn, [],
+        break_point,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=False)
+
+    actual = []
+    # Restore from checkpoint and then run init_op.
+    with ops.Graph().as_default() as g:
+      saver = self._import_meta_graph()
+      init_op, get_next_op = self._get_iterator_ops_from_collection(
+          ds_fn, sparse_tensors=sparse_tensors)
+      get_next_op = remove_variants(get_next_op)
+      with self.session(graph=g) as sess:
+        self._initialize(init_op, sess)
+        self._restore(saver, sess)
+        self._initialize(init_op, sess)
+        for _ in range(num_outputs):
+          actual.append(sess.run(get_next_op))
+        if verify_exhausted:
+          with self.assertRaises(errors.OutOfRangeError):
+            sess.run(get_next_op)
+    self.match(expected, actual, assert_items_equal=assert_items_equal)
+
+  def verify_error_on_save(self,
+                           ds_fn,
+                           num_outputs,
+                           error,
+                           break_point=None,
+                           sparse_tensors=False,
+                           assert_items_equal=False):
+    """Attempts to save a non-saveable iterator.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      num_outputs: Total number of outputs expected from this Dataset.
+      error: Declared error when trying to save iterator.
+      break_point: Break point. Optional. Defaults to num_outputs/2.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    del assert_items_equal
+    break_point = num_outputs // 2 if not break_point else break_point
+    if context.executing_eagerly():
+      iterator = iter(ds_fn())
+      ckpt = tracking_util.Checkpoint(iterator=iterator)
+      for _ in range(break_point):
+        next(iterator)
+      with self.assertRaises(error):
+        ckpt.save(self._ckpt_path())
+    else:
+      with ops.Graph().as_default() as g:
+        init_op, get_next_op, saver = self._build_graph(
+            ds_fn, sparse_tensors=sparse_tensors)
+        get_next_op = remove_variants(get_next_op)
+        with self.session(graph=g) as sess:
+          self._initialize(init_op, sess)
+          for _ in range(break_point):
+            sess.run(get_next_op)
+          with self.assertRaises(error):
+            self._save(sess, saver)
+
+  def verify_run_with_breaks(self,
+                             ds_fn,
+                             break_points,
+                             num_outputs,
+                             sparse_tensors=False,
+                             verify_exhausted=True,
+                             assert_items_equal=False):
+    """Verifies that ds_fn() produces the same outputs with and without breaks.
+
+    1. Builds a Dataset using `ds_fn` and produces `num_outputs` items from it
+       *without* stopping at break points.
+    2. Builds a Dataset using `ds_fn` and produces `num_outputs` items from it
+       with stopping at break points.
+
+    Deep matches outputs from 1 and 2.
+
+    Args:
+      ds_fn: 0-argument function that returns a Dataset.
+      break_points: A list of integers. For each `break_point` in
+        `break_points`, we produce outputs till `break_point` number of items
+        have been produced and then checkpoint the state. The current graph and
+        session are destroyed and a new graph and session are used to produce
+        outputs till next checkpoint or till `num_outputs` elements have been
+        produced. `break_point` must be <= `num_outputs`.
+      num_outputs: Total number of outputs expected from this Dataset.
+      sparse_tensors: Whether dataset is built from SparseTensor(s).
+      verify_exhausted: Whether to verify that the iterator has been exhausted
+        after producing `num_outputs` elements.
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if any test fails.
+    """
+    expected = self.gen_outputs(
+        ds_fn, [],
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=verify_exhausted)
+
+    actual = self.gen_outputs(
+        ds_fn,
+        break_points,
+        num_outputs,
+        sparse_tensors=sparse_tensors,
+        verify_exhausted=verify_exhausted)
+
+    self.match(expected, actual, assert_items_equal=assert_items_equal)
+
+  def gen_outputs(self,
+                  ds_fn,
+                  break_points,
+                  num_outputs,
+                  ckpt_saved=False,
+                  sparse_tensors=False,
+                  verify_exhausted=True,
+                  save_checkpoint_at_end=True):
+    """Generates elements from input dataset while stopping at break points.
+
+    Produces `num_outputs` outputs and saves the state of the iterator in the
+    Saver checkpoint.
+
+    Args:
+      ds_fn: 0-argument function that returns the dataset.
+      break_points: A list of integers. For each `break_point` in
+        `break_points`, we produce outputs till `break_point` number of items
+        have been produced and then checkpoint the state. The current graph and
+        session are destroyed and a new graph and session are used to produce
+        outputs till next checkpoint or till `num_outputs` elements have been
+        produced. `break_point` must be <= `num_outputs`.
+      num_outputs: The total number of outputs to produce from the iterator.
+      ckpt_saved: Whether a checkpoint already exists.
+      sparse_tensors:  Whether dataset is built from SparseTensor(s).
+      verify_exhausted: Whether to verify that the iterator has been exhausted
+        after producing `num_outputs` elements.
+      save_checkpoint_at_end: Whether to save a checkpoint after producing all
+        outputs. If False, checkpoints are saved each break point but not at the
+        end. Note that checkpoints overwrite each other so there is always only
+        a single checkpoint available. Defaults to True.
+
+    Returns:
+      A list of `num_outputs` items.
+    """
+    outputs = []
+
+    if context.executing_eagerly():
+      for i in range(len(break_points) + 1):
+        iterator = iter(ds_fn())
+        ckpt = tracking_util.Checkpoint(iterator=iterator)
+        if ckpt_saved:
+          ckpt_path = self._latest_ckpt()
+          ckpt.restore(ckpt_path)
+        start = break_points[i - 1] if i > 0 else 0
+        end = break_points[i] if i < len(break_points) else num_outputs
+        num_iters = end - start
+        for _ in range(num_iters):
+          outputs.append(self.evaluate(next(iterator)))
+        if i < len(break_points) and end == num_outputs and verify_exhausted:
+          # If the checkpoint is expected to be after the final element, makes
+          # sure the iterator is exhausted. Otherwise, it may happen that the
+          # iterator has produced `num_outputs` elements, but the source dataset
+          # has more elements, resulting in saving incorrect element counts in
+          # some source datasets.
+          # For example: `range(10).shuffle(10).filter(lambda x: x % 2 == 0)`
+          # Calling `next` one more time exhausts all upstream iterators.
+          with self.assertRaises(StopIteration):
+            next(iterator)
+        if i == len(break_points) and verify_exhausted:
+          with self.assertRaises(StopIteration):
+            next(iterator)
+        if save_checkpoint_at_end or i < len(break_points):
+          # TODO(b/275117275): Verify if TF2 async checkpoint works.
+          ckpt_options = checkpoint_options.CheckpointOptions()
+          ckpt_options.experimental_enable_async_checkpoint = False
+          ckpt_options.enable_async = False
+          ckpt_path = ckpt.save(self._ckpt_path(), options=ckpt_options)
+          ckpt_saved = True
+    else:
+      def get_ops():
+        if ckpt_saved:
+          saver = self._import_meta_graph()
+          init_op, get_next_op = self._get_iterator_ops_from_collection(
+              ds_fn, sparse_tensors=sparse_tensors)
+        else:
+          init_op, get_next_op, saver = self._build_graph(
+              ds_fn, sparse_tensors=sparse_tensors)
+        return init_op, get_next_op, saver
+
+      for i in range(len(break_points) + 1):
+        with ops.Graph().as_default() as g:
+          init_op, get_next_op, saver = get_ops()
+          get_next_op = remove_variants(get_next_op)
+          with self.session(graph=g) as sess:
+            if ckpt_saved:
+              self._initialize(init_op, sess)
+              self._restore(saver, sess)
+            else:
+              self._initialize(init_op, sess)
+            start = break_points[i - 1] if i > 0 else 0
+            end = break_points[i] if i < len(break_points) else num_outputs
+            num_iters = end - start
+            for _ in range(num_iters):
+              outputs.append(sess.run(get_next_op))
+            if (i < len(break_points) and end == num_outputs and
+                verify_exhausted):
+              with self.assertRaises(errors.OutOfRangeError):
+                sess.run(get_next_op)
+            if i == len(break_points) and verify_exhausted:
+              with self.assertRaises(errors.OutOfRangeError):
+                sess.run(get_next_op)
+            if save_checkpoint_at_end or i < len(break_points):
+              self._save(sess, saver)
+              ckpt_saved = True
+
+    return outputs
+
+  def match(self, expected, actual, assert_items_equal=False):
+    """Matches nested structures.
+
+    Recursively matches shape and values of `expected` and `actual`.
+    Handles scalars, numpy arrays and other python sequence containers
+    e.g. list, dict, as well as SparseTensorValue and RaggedTensorValue.
+
+    Args:
+      expected: Nested structure 1.
+      actual: Nested structure 2.
+      assert_items_equal: Tests the output has the expected elements regardless
+        of order.
+
+    Raises:
+      AssertionError if matching fails.
+    """
+    if isinstance(expected, np.ndarray):
+      expected = expected.tolist()
+    if isinstance(actual, np.ndarray):
+      actual = actual.tolist()
+    self.assertEqual(type(expected), type(actual))
+
+    if nest.is_nested(expected):
+      self.assertEqual(len(expected), len(actual))
+      if isinstance(expected, dict):
+        for key1, key2 in zip(sorted(expected), sorted(actual)):
+          self.assertEqual(key1, key2)
+          self.match(expected[key1], actual[key2])
+      elif assert_items_equal:
+        for item1, item2 in zip(sorted(expected), sorted(actual)):
+          self.match(item1, item2)
+      else:
+        for item1, item2 in zip(expected, actual):
+          self.match(item1, item2)
+    elif isinstance(expected, sparse_tensor.SparseTensorValue):
+      self.match((expected.indices, expected.values, expected.dense_shape),
+                 (actual.indices, actual.values, actual.dense_shape))
+    elif isinstance(expected, ragged_tensor_value.RaggedTensorValue):
+      self.match((expected.values, expected.row_splits),
+                 (actual.values, actual.row_splits))
+    else:
+      self.assertEqual(expected, actual)
+
+  def does_not_match(self, expected, actual):
+    with self.assertRaises(AssertionError):
+      self.match(expected, actual)
+
+  def gen_break_points(self, num_outputs, num_samples=10):
+    """Generates `num_samples` unique break points in [0, num_outputs]."""
+    return np.unique(np.linspace(0, num_outputs, num_samples, dtype=int))
+
+  def _build_graph(self, ds_fn, sparse_tensors=False):
+    dataset = ds_fn()
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    external_state_policy = dataset.options().experimental_external_state_policy
+    saveable = contrib_iterator_ops.make_saveable_from_iterator(
+        iterator, external_state_policy=external_state_policy)
+    ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable)
+    init_op = iterator.initializer
+    if sparse_tensors:
+      get_next = sparse_tensor.SparseTensor(*iterator.get_next())
+    else:
+      get_next = iterator.get_next()
+    self._add_iterator_ops_to_collection(init_op, get_next, ds_fn,
+                                         sparse_tensors)
+    saver = saver_lib.Saver(allow_empty=True)
+    return init_op, get_next, saver
+
+  def _add_iterator_ops_to_collection(self,
+                                      init_op,
+                                      get_next,
+                                      ds_fn,
+                                      sparse_tensors=False):
+    ops.add_to_collection("iterator_ops", init_op)
+    # `get_next` may be a tuple e.g. in TensorSliceDataset. Since Collections
+    # do not support tuples we flatten the tensors and restore the shape in
+    # `_get_iterator_ops_from_collection`.
+    if sparse_tensors:  # specific for deprecated `from_sparse_tensor_slices`.
+      ops.add_to_collection("iterator_ops", get_next.indices)
+      ops.add_to_collection("iterator_ops", get_next.values)
+      ops.add_to_collection("iterator_ops", get_next.dense_shape)
+      return
+
+    get_next_list = nest.flatten(get_next)
+    for i, output_class in enumerate(
+        nest.flatten(self._get_output_classes(ds_fn))):
+      if output_class is sparse_tensor.SparseTensor:
+        ops.add_to_collection("iterator_ops", get_next_list[i].indices)
+        ops.add_to_collection("iterator_ops", get_next_list[i].values)
+        ops.add_to_collection("iterator_ops", get_next_list[i].dense_shape)
+      else:
+        ops.add_to_collection("iterator_ops", get_next_list[i])
+
+  def _get_iterator_ops_from_collection(self, ds_fn, sparse_tensors=False):
+    all_ops = ops.get_collection("iterator_ops")
+    if sparse_tensors:  # specific for deprecated `from_sparse_tensor_slices`.
+      init_op, indices, values, dense_shape = all_ops
+      return init_op, sparse_tensor.SparseTensor(indices, values, dense_shape)
+    get_next_list = []
+    i = 1
+    for output_class in nest.flatten(self._get_output_classes(ds_fn)):
+      if output_class is sparse_tensor.SparseTensor:
+        indices, values, dense_shape = all_ops[i:i + 3]
+        i += 3
+        get_next_list.append(
+            sparse_tensor.SparseTensor(indices, values, dense_shape))
+      else:
+        get_next_list.append(all_ops[i])
+        i += 1
+    return all_ops[0], nest.pack_sequence_as(
+        self._get_output_types(ds_fn), get_next_list)
+
+  def _get_output_types(self, ds_fn):
+    assert not context.executing_eagerly()
+    with ops.Graph().as_default():
+      return dataset_ops.get_legacy_output_types(ds_fn())
+
+  def _get_output_shapes(self, ds_fn):
+    assert not context.executing_eagerly()
+    with ops.Graph().as_default():
+      return dataset_ops.get_legacy_output_shapes(ds_fn())
+
+  def _get_output_classes(self, ds_fn):
+    assert not context.executing_eagerly()
+    with ops.Graph().as_default():
+      return dataset_ops.get_legacy_output_classes(ds_fn())
+
+  def _ckpt_path(self):
+    return os.path.join(self.get_temp_dir(), "iterator")
+
+  def _latest_ckpt(self):
+    return checkpoint_management.latest_checkpoint(self.get_temp_dir())
+
+  def _save(self, sess, saver):
+    saver.save(sess, self._ckpt_path())
+
+  def _restore(self, saver, sess):
+    sess.run(lookup_ops.tables_initializer())
+    saver.restore(sess, self._latest_ckpt())
+
+  def _initialize(self, init_op, sess):
+    sess.run(variables.global_variables_initializer())
+    sess.run(lookup_ops.tables_initializer())
+    sess.run(init_op)
+
+  def _import_meta_graph(self):
+    meta_file_path = self._ckpt_path() + ".meta"
+    return saver_lib.import_meta_graph(meta_file_path)
+
+  def _delete_ckpt(self):
+    # Remove all checkpoint files.
+    prefix = self._ckpt_path()
+    pattern = prefix + "*"
+    files = gfile.Glob(pattern)
+    map(gfile.Remove, files)

tensorflow/python/data/kernel_tests/choose_from_datasets_test.py

@@ -0,0 +1,172 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.choose_from_dataset()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.platform import test
+
+
+class ChooseFromDatasetsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasets(self):
+    words = [b"foo", b"bar", b"baz"]
+    datasets = [dataset_ops.Dataset.from_tensors(w).repeat() for w in words]
+    choice_array = np.random.randint(3, size=(15,), dtype=np.int64)
+    choice_dataset = dataset_ops.Dataset.from_tensor_slices(choice_array)
+    dataset = dataset_ops.Dataset.choose_from_datasets(datasets, choice_dataset)
+    next_element = self.getNext(dataset)
+    for i in choice_array:
+      self.assertEqual(words[i], self.evaluate(next_element()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasetsStoppingOnEmptyDataset(self):
+    datasets = [
+        dataset_ops.Dataset.from_tensors(b"foo").repeat(2),
+        dataset_ops.Dataset.from_tensors(b"bar").repeat(),
+        dataset_ops.Dataset.from_tensors(b"baz").repeat(),
+    ]
+    choice_array = np.asarray([0, 0, 0, 1, 1, 1, 2, 2, 2], dtype=np.int64)
+    choice_dataset = dataset_ops.Dataset.from_tensor_slices(choice_array)
+    dataset = dataset_ops.Dataset.choose_from_datasets(
+        datasets, choice_dataset, stop_on_empty_dataset=True)
+    self.assertDatasetProduces(dataset, [b"foo", b"foo"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasetsSkippingEmptyDatasets(self):
+    datasets = [
+        dataset_ops.Dataset.from_tensors(b"foo").repeat(2),
+        dataset_ops.Dataset.from_tensors(b"bar").repeat(),
+        dataset_ops.Dataset.from_tensors(b"baz").repeat(),
+    ]
+    choice_array = np.asarray([0, 0, 0, 1, 1, 1, 2, 2, 2], dtype=np.int64)
+    choice_dataset = dataset_ops.Dataset.from_tensor_slices(choice_array)
+    dataset = dataset_ops.Dataset.choose_from_datasets(
+        datasets, choice_dataset, stop_on_empty_dataset=False)
+    # Chooses 2 elements from the first dataset while the selector specifies 3.
+    self.assertDatasetProduces(
+        dataset,
+        [b"foo", b"foo", b"bar", b"bar", b"bar", b"baz", b"baz", b"baz"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasetsChoiceDatasetIsEmpty(self):
+    datasets = [
+        dataset_ops.Dataset.from_tensors(b"foo").repeat(),
+        dataset_ops.Dataset.from_tensors(b"bar").repeat(),
+        dataset_ops.Dataset.from_tensors(b"baz").repeat(),
+    ]
+    dataset = dataset_ops.Dataset.choose_from_datasets(
+        datasets,
+        choice_dataset=dataset_ops.Dataset.range(0),
+        stop_on_empty_dataset=False)
+    self.assertDatasetProduces(dataset, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasetsNested(self):
+    ds1 = dataset_ops.Dataset.range(10).window(2)
+    ds2 = dataset_ops.Dataset.range(10, 20).window(2)
+    choice_dataset = dataset_ops.Dataset.range(2).repeat(5)
+    ds = dataset_ops.Dataset.choose_from_datasets([ds1, ds2], choice_dataset)
+    ds = ds.flat_map(lambda x: x)
+    expected = []
+    for i in range(5):
+      for j in range(2):
+        expected.extend([10*j + 2*i, 10*j + 2*i + 1])
+    self.assertDatasetProduces(ds, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrors(self):
+    with self.assertRaisesRegex(TypeError, "tf.int64"):
+      dataset_ops.Dataset.choose_from_datasets(
+          [
+              dataset_ops.Dataset.from_tensors(0),
+              dataset_ops.Dataset.from_tensors(1)
+          ],
+          choice_dataset=dataset_ops.Dataset.from_tensors(1.0))
+
+    with self.assertRaisesRegex(TypeError, "scalar"):
+      dataset_ops.Dataset.choose_from_datasets(
+          [
+              dataset_ops.Dataset.from_tensors(0),
+              dataset_ops.Dataset.from_tensors(1)
+          ],
+          choice_dataset=dataset_ops.Dataset.from_tensors([1.0]))
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError, "out of range"):
+      dataset = dataset_ops.Dataset.choose_from_datasets(
+          [dataset_ops.Dataset.from_tensors(0)],
+          choice_dataset=dataset_ops.Dataset.from_tensors(
+              constant_op.constant(1, dtype=dtypes.int64)))
+      next_element = self.getNext(dataset)
+      self.evaluate(next_element())
+
+    with self.assertRaisesRegex(
+        ValueError, r"Invalid `datasets`. `datasets` should not be empty."):
+      dataset_ops.Dataset.choose_from_datasets(
+          datasets=[], choice_dataset=dataset_ops.Dataset.from_tensors(1.0))
+
+    with self.assertRaisesRegex(
+        TypeError, r"`choice_dataset` should be a `tf.data.Dataset`"):
+      datasets = [dataset_ops.Dataset.range(42)]
+      dataset_ops.Dataset.choose_from_datasets(datasets, choice_dataset=None)
+
+
+class ChooseFromDatasetsCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                       parameterized.TestCase):
+
+  def _build_dataset(self,
+                     num_datasets,
+                     num_elements_per_dataset,
+                     options=None):
+    datasets = [
+        dataset_ops.Dataset.range(num_elements_per_dataset)
+        for _ in range(num_datasets)
+    ]
+    indices = []
+    for i in range(num_datasets):
+      indices = indices + ([i] * num_elements_per_dataset)
+    shuffled_indices = stateless_random_ops.stateless_shuffle(
+        np.int64(indices), seed=[1, 2])
+    choice_dataset = dataset_ops.Dataset.from_tensor_slices(shuffled_indices)
+    dataset = dataset_ops.Dataset.choose_from_datasets(datasets, choice_dataset)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self, lambda: self._build_dataset(5, 20, options), num_outputs=100)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/concatenate_test.py

@@ -0,0 +1,550 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.concatenate()."""
+from typing import Callable, Tuple
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+
+
+class ConcatenateTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    input_components = (
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 15),
+        np.array([37.0, 38.0, 39.0, 40.0]))
+    to_concatenate_components = (
+        np.tile(np.array([[1], [2], [3], [4], [5]]), 20),
+        np.tile(np.array([[12], [13], [14], [15], [16]]), 15),
+        np.array([37.0, 38.0, 39.0, 40.0, 41.0]))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components)
+    dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices(
+        to_concatenate_components)
+    concatenated = input_dataset.concatenate(dataset_to_concatenate)
+    self.assertEqual(
+        dataset_ops.get_legacy_output_shapes(concatenated),
+        (tensor_shape.TensorShape([20]), tensor_shape.TensorShape([15]),
+         tensor_shape.TensorShape([])))
+
+    get_next = self.getNext(concatenated)
+
+    for i in range(9):
+      result = self.evaluate(get_next())
+      if i < 4:
+        for component, result_component in zip(input_components, result):
+          self.assertAllEqual(component[i], result_component)
+      else:
+        for component, result_component in zip(to_concatenate_components,
+                                               result):
+          self.assertAllEqual(component[i - 4], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentShape(self):
+    input_components = (
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 4))
+    to_concatenate_components = (
+        np.tile(np.array([[1], [2], [3], [4], [5]]), 20),
+        np.tile(np.array([[12], [13], [14], [15], [16]]), 15))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components)
+    dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices(
+        to_concatenate_components)
+    concatenated = input_dataset.concatenate(dataset_to_concatenate)
+    self.assertEqual(
+        [ts.as_list()
+         for ts in nest.flatten(
+             dataset_ops.get_legacy_output_shapes(concatenated))],
+        [[20], [None]])
+    get_next = self.getNext(concatenated)
+    for i in range(9):
+      result = self.evaluate(get_next())
+      if i < 4:
+        for component, result_component in zip(input_components, result):
+          self.assertAllEqual(component[i], result_component)
+      else:
+        for component, result_component in zip(to_concatenate_components,
+                                               result):
+          self.assertAllEqual(component[i - 4], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentStructure(self):
+    input_components = (
+        np.tile(np.array([[1], [2], [3], [4]]), 5),
+        np.tile(np.array([[12], [13], [14], [15]]), 4))
+    to_concatenate_components = (
+        np.tile(np.array([[1], [2], [3], [4], [5]]), 20),
+        np.tile(np.array([[12], [13], [14], [15], [16]]), 15),
+        np.array([37.0, 38.0, 39.0, 40.0, 41.0]))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components)
+    dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices(
+        to_concatenate_components)
+
+    with self.assertRaisesRegex(TypeError, "Incompatible dataset elements"):
+      input_dataset.concatenate(dataset_to_concatenate)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentKeys(self):
+    input_components = {
+        "foo": np.array([[1], [2], [3], [4]]),
+        "bar": np.array([[12], [13], [14], [15]])
+    }
+    to_concatenate_components = {
+        "foo": np.array([[1], [2], [3], [4]]),
+        "baz": np.array([[5], [6], [7], [8]])
+    }
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components)
+    dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices(
+        to_concatenate_components)
+
+    with self.assertRaisesRegex(TypeError, "Incompatible dataset elements"):
+      input_dataset.concatenate(dataset_to_concatenate)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentType(self):
+    input_components = (
+        np.tile(np.array([[1], [2], [3], [4]]), 5),
+        np.tile(np.array([[12], [13], [14], [15]]), 4))
+    to_concatenate_components = (
+        np.tile(np.array([[1.0], [2.0], [3.0], [4.0]]), 5),
+        np.tile(np.array([[12], [13], [14], [15]]), 15))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(input_components)
+    dataset_to_concatenate = dataset_ops.Dataset.from_tensor_slices(
+        to_concatenate_components)
+
+    with self.assertRaisesRegex(TypeError, "Incompatible dataset elements"):
+      input_dataset.concatenate(dataset_to_concatenate)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindows(self):
+    a = dataset_ops.Dataset.range(5).window(1)
+    b = dataset_ops.Dataset.range(5, 10).window(1)
+    c = a.concatenate(b).flat_map(lambda x: x)
+    self.assertDatasetProduces(c, list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    a = dataset_ops.Dataset.range(5)
+    b = dataset_ops.Dataset.range(5, 10)
+    c = a.concatenate(b, name="concatenate")
+    self.assertDatasetProduces(c, list(range(10)))
+
+
+class ConcatenateCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  def _build_concatenate_dataset(self, var_array, options=None):
+    input_components = (np.tile(np.array([[1], [2], [3], [4]]), 20),
+                        np.tile(np.array([[12], [13], [14], [15]]), 4))
+    to_concatenate_components = (np.tile(
+        np.array([[5], [6], [7], [8], [9]]), 20), var_array)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        input_components).concatenate(
+            dataset_ops.Dataset.from_tensor_slices(to_concatenate_components))
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    num_outputs = 9
+    array = np.tile(np.array([[16], [17], [18], [19], [20]]), 15)
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self._build_concatenate_dataset(array, options),
+              num_outputs)
+
+
+class ConcatenateRandomAccessTest(test_base.DatasetTestBase,
+                                  parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 3, 4])))
+  def testInvalidIndex(self, index):
+    input_dataset = dataset_ops.Dataset.from_tensor_slices([-1])
+    concatenate_dataset = dataset_ops.Dataset.from_tensor_slices([1, 2])
+    concatenated = input_dataset.concatenate(concatenate_dataset)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(concatenated, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testConcatenateTwoEmptyDatasets(self):
+    input_dataset = dataset_ops.Dataset.from_tensor_slices([])
+    concatenate_dataset = dataset_ops.Dataset.from_tensor_slices([])
+    concatenated = input_dataset.concatenate(concatenate_dataset)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(concatenated, index=0))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testConcatenateAnEmptyDataset(self):
+    input_dataset = dataset_ops.Dataset.from_tensor_slices([1.0])
+    concatenate_dataset = dataset_ops.Dataset.from_tensor_slices([])
+    concatenated = input_dataset.concatenate(concatenate_dataset)
+    self.assertAllEqual(
+        self.evaluate(random_access.at(concatenated, index=0)), 1.0)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(concatenated, index=1))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testConcatenateOntoEmptyDataset(self):
+    input_dataset = dataset_ops.Dataset.from_tensor_slices([])
+    concatenate_dataset = dataset_ops.Dataset.from_tensor_slices([2.0, 3.0])
+    concatenated = input_dataset.concatenate(concatenate_dataset)
+    self.assertAllEqual(
+        self.evaluate(random_access.at(concatenated, index=0)), 2.0)
+    self.assertAllEqual(
+        self.evaluate(random_access.at(concatenated, index=1)), 3.0)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(concatenated, index=2))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testConcatenateTwoNonEmptyDatasets(self):
+    input_dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2])
+    concatenate_dataset = dataset_ops.Dataset.from_tensor_slices([3, 4])
+    concatenated = input_dataset.concatenate(concatenate_dataset)
+    for i in range(5):
+      self.assertAllEqual(random_access.at(concatenated, index=i), i)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(concatenated, index=5))
+
+
+class GlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for global shuffling of tf.data datasets."""
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffledOutput(self):
+    dataset1 = dataset_ops.Dataset.range(0, 5)
+    dataset2 = dataset_ops.Dataset.range(5, 17)
+
+    dataset = dataset1.concatenate(dataset2)
+
+    dataset = global_shuffle_op._global_shuffle(dataset)
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(output, range(0, 17))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffledWithBatchOutput(self):
+    """Testing with `.batch()` ensures the global shuffle map is stateless."""
+    dataset1 = dataset_ops.Dataset.range(0, 4)
+    dataset2 = dataset_ops.Dataset.range(4, 10)
+
+    dataset = dataset1.concatenate(dataset2)
+    dataset = dataset.batch(3, drop_remainder=True)
+
+    dataset = global_shuffle_op._global_shuffle(dataset)
+
+    got = self.getDatasetOutput(dataset, requires_initialization=True)
+    expected = [
+        np.array([0, 1, 2], dtype=np.int32),
+        np.array([3, 4, 5], dtype=np.int32),
+        np.array([6, 7, 8], dtype=np.int32),
+    ]
+
+    self.assertIsInstance(got, list)
+    # Converts to tuples for lexicographically sort
+    got.sort(key=tuple)
+
+    self.assertLen(got, len(expected))
+
+    for element_got, element_expected in zip(got, expected):
+      self.assertAllEqual(element_got, element_expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedConcatenateShuffledOutput(self):
+    dataset1 = dataset_ops.Dataset.range(0, 3)
+    dataset2 = dataset_ops.Dataset.range(3, 6)
+    dataset3 = dataset_ops.Dataset.range(6, 9)
+
+    dataset = dataset1.concatenate(dataset2)
+    dataset = dataset.concatenate(dataset3)
+
+    dataset = global_shuffle_op._global_shuffle(dataset)
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(output, range(0, 9))
+
+
+class ConcatenateGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(10, 8), (9, 5), (4, 7), (5, 8)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testConcatenate(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, int],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+
+    def _build_dataset():
+      first_dataset = dataset_ops.Dataset.range(dataset_ranges[0])
+      second_dataset = dataset_ops.Dataset.range(
+          dataset_ranges[0], dataset_ranges[0] + dataset_ranges[1]
+      )
+      dataset = first_dataset.concatenate(second_dataset)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=sum(dataset_ranges),
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(10, 8, 11), (9, 5, 3)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testNestedConcatenate(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, int],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+
+    def _build_dataset():
+      first_dataset = dataset_ops.Dataset.range(dataset_ranges[0])
+      second_dataset = dataset_ops.Dataset.range(
+          dataset_ranges[0], dataset_ranges[0] + dataset_ranges[1]
+      )
+      third_dataset = dataset_ops.Dataset.range(
+          sum(dataset_ranges[:2]), sum(dataset_ranges[:3])
+      )
+
+      dataset = first_dataset.concatenate(second_dataset)
+      dataset = dataset.concatenate(third_dataset)
+
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=sum(dataset_ranges),
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(3, 4, 6, 5)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testFourNestedConcatenate(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, int],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+    def _build_dataset():
+      first_dataset = dataset_ops.Dataset.range(dataset_ranges[0])
+      second_dataset = dataset_ops.Dataset.range(
+          dataset_ranges[0], sum(dataset_ranges[:2])
+      )
+      third_dataset = dataset_ops.Dataset.range(
+          sum(dataset_ranges[:2]), sum(dataset_ranges[:3])
+      )
+      fourth_dataset = dataset_ops.Dataset.range(
+          sum(dataset_ranges[:3]), sum(dataset_ranges)
+      )
+
+      left = first_dataset.concatenate(second_dataset)
+      right = third_dataset.concatenate(fourth_dataset)
+
+      dataset = left.concatenate(right)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=sum(dataset_ranges),
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(1, 2, 3, 4, 5, 6)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testDeepConcatenate(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, ...],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+    def _build_dataset():
+      prefix_sums = [0] * (len(dataset_ranges) + 1)
+      for i, value in enumerate(dataset_ranges):
+        prefix_sums[i + 1] = prefix_sums[i] + value
+
+      dataset = dataset_ops.Dataset.range(prefix_sums[0], prefix_sums[1])
+      for i in range(1, len(dataset_ranges)):
+        to_concat = dataset_ops.Dataset.range(
+            prefix_sums[i], prefix_sums[i + 1]
+        )
+        dataset = dataset.concatenate(to_concat)
+
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=sum(dataset_ranges),
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(1, 2, 3, 4, 5, 6)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testDeepConcatenateWithBatchAndPrefetch(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, ...],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+    def _build_dataset():
+      prefix_sums = [0] * (len(dataset_ranges) + 1)
+      for i, value in enumerate(dataset_ranges):
+        prefix_sums[i + 1] = prefix_sums[i] + value
+
+      dataset = dataset_ops.Dataset.range(prefix_sums[0], prefix_sums[1])
+      for i in range(1, len(dataset_ranges)):
+        to_concat = dataset_ops.Dataset.range(
+            prefix_sums[i], prefix_sums[i + 1]
+        )
+        dataset = dataset.concatenate(to_concat)
+
+      dataset = dataset.batch(2, drop_remainder=True)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+      dataset = dataset.unbatch()
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=(sum(dataset_ranges) // 2) * 2,
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/counter_test.py

@@ -0,0 +1,72 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.counter`."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+
+
+class CounterTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(start=3, step=4, expected_output=[[3, 7, 11]]) +
+          combinations.combine(start=0, step=-1, expected_output=[[0, -1, -2]]))
+  )
+  def testCounter(self, start, step, expected_output):
+    dataset = dataset_ops.Dataset.counter(start, step)
+    self.assertEqual(
+        [], dataset_ops.get_legacy_output_shapes(dataset).as_list())
+    self.assertEqual(dtypes.int64, dataset_ops.get_legacy_output_types(dataset))
+    get_next = self.getNext(dataset)
+    for expected in expected_output:
+      self.assertEqual(expected, self.evaluate(get_next()))
+
+
+class CounterCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                            parameterized.TestCase):
+
+  def _build_counter_dataset(self, start, step, num_outputs, options=None):
+    counter_dataset = dataset_ops.Dataset.counter(start, step)
+    range_dataset = dataset_ops.Dataset.range(num_outputs)
+    dataset = dataset_ops.Dataset.zip((counter_dataset, range_dataset))
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    num_outputs = 10
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self, lambda: self._build_counter_dataset(
+            start=2, step=10, num_outputs=num_outputs, options=options),
+        num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/dataset_spec_test.py

@@ -0,0 +1,136 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.DatasetSpec`."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.platform import test
+
+
+class DatasetSpecTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInputSignature(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        np.arange(10).astype(np.int32)).batch(5)
+
+    @def_function.function(input_signature=[
+        dataset_ops.DatasetSpec(
+            tensor_spec.TensorSpec(
+                shape=(None,), dtype=dtypes.int32, name=None),
+            tensor_shape.TensorShape([]))
+    ])
+    def fn(_):
+      pass
+
+    fn(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetSpecInnerSpec(self):
+    inner_spec = tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    ds_spec = dataset_ops.DatasetSpec(inner_spec)
+    self.assertEqual(ds_spec.element_spec, inner_spec)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetSpecTraceType(self):
+    trace_type_1 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32),
+        [5])
+    trace_type_2 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32),
+        [5])
+
+    self.assertEqual(trace_type_1, trace_type_2)
+    self.assertEqual(hash(trace_type_1), hash(trace_type_2))
+    self.assertTrue(trace_type_1.is_subtype_of(trace_type_2))
+    self.assertTrue(trace_type_2.is_subtype_of(trace_type_1))
+
+    trace_type_3 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32),
+        [6])
+    self.assertNotEqual(trace_type_1, trace_type_3)
+    self.assertFalse(trace_type_1.is_subtype_of(trace_type_3))
+    self.assertFalse(trace_type_3.is_subtype_of(trace_type_1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetSpecHierarchical(self):
+    spec_1 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(1, None), dtype=dtypes.int32),
+        [5, None, 2])
+    spec_2 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(None, None), dtype=dtypes.int32),
+        [None, None, None])
+    spec_3 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(1, 2), dtype=dtypes.int32),
+        [5, 3, 2])
+    spec_4 = dataset_ops.DatasetSpec(
+        tensor_spec.TensorSpec(shape=(None, 2), dtype=dtypes.int32),
+        [None, 1, None])
+
+    self.assertTrue(spec_1.is_subtype_of(spec_1))
+
+    self.assertTrue(spec_1.is_subtype_of(spec_2))
+    self.assertTrue(spec_3.is_subtype_of(spec_2))
+    self.assertTrue(spec_4.is_subtype_of(spec_2))
+
+    self.assertFalse(spec_2.is_subtype_of(spec_1))
+    self.assertFalse(spec_2.is_subtype_of(spec_3))
+    self.assertFalse(spec_2.is_subtype_of(spec_4))
+
+    self.assertEqual(spec_1.most_specific_common_supertype([]), spec_1)
+    self.assertEqual(spec_1.most_specific_common_supertype([spec_4]), spec_2)
+    self.assertEqual(
+        spec_1.most_specific_common_supertype([spec_3, spec_4]), spec_2)
+    self.assertEqual(
+        spec_1.most_specific_common_supertype([spec_2, spec_3, spec_4]), spec_2)
+
+  # TODO(b/220385675): element_spec should always be a TypeSpec.
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetSpecHierarchicalDict(self):
+    spec_1 = dataset_ops.DatasetSpec(
+        {"a": tensor_spec.TensorSpec(shape=(1, None), dtype=dtypes.int32)},
+        [])
+    spec_2 = dataset_ops.DatasetSpec(
+        {"a": tensor_spec.TensorSpec(shape=(None, None), dtype=dtypes.int32)},
+        [])
+    spec_3 = dataset_ops.DatasetSpec(
+        {"b": tensor_spec.TensorSpec(shape=(1, None), dtype=dtypes.int32)},
+        [])
+    spec_4 = dataset_ops.DatasetSpec({"b": None}, [])
+
+    self.assertTrue(spec_1.is_subtype_of(spec_1))
+    self.assertTrue(spec_1.is_subtype_of(spec_2))
+    self.assertFalse(spec_2.is_subtype_of(spec_1))
+
+    self.assertFalse(spec_1.is_subtype_of(spec_3))
+    self.assertFalse(spec_3.is_subtype_of(spec_1))
+    self.assertFalse(spec_2.is_subtype_of(spec_3))
+    self.assertFalse(spec_3.is_subtype_of(spec_2))
+
+    self.assertTrue(spec_4.is_subtype_of(spec_4))
+    self.assertEqual(spec_4.most_specific_common_supertype([]), spec_4)
+    self.assertEqual(spec_4.most_specific_common_supertype([spec_4]), spec_4)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/dataset_test.py

@@ -0,0 +1,692 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset`."""
+
+import collections
+import os
+import warnings
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.data.ops import optional_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import readers
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.lib.io import tf_record
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+
+class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAsSerializedGraph(self):
+    dataset = dataset_ops.Dataset.range(10)
+    graph = graph_pb2.GraphDef().FromString(
+        self.evaluate(dataset._as_serialized_graph()))
+    self.assertTrue(any(node.op == "RangeDataset" for node in graph.node))
+
+  def testAsSerializedGraphStateful(self):
+    dataset = dataset_ops.Dataset.range(10).map(
+        lambda _: random_ops.random_uniform(()))
+    with self.assertRaises(errors.FailedPreconditionError):
+      self.evaluate(
+          dataset._as_serialized_graph(external_state_policy=options_lib
+                                       .ExternalStatePolicy.FAIL))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              init_source=["textfile", "keyvaluetensor", "dataset"])))
+  def testLookupTableGraphSerialization(self, init_source):
+    vals = [10, 11]
+    initializer = self.lookupTableInitializer(init_source, vals)
+    table = lookup_ops.StaticHashTable(initializer, -1)
+    dataset = dataset_ops.Dataset.range(3)
+    dataset = dataset.map(table.lookup)
+    self.evaluate(lookup_ops.tables_initializer())
+    round_tripped = self.graphRoundTrip(dataset)
+    del table
+    del dataset
+    self.assertDatasetProduces(
+        round_tripped, [10, 11, -1], requires_initialization=True)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testAsFunctionWithMap(self):
+    with ops.device("CPU"):
+      original_dataset = dataset_ops.Dataset.range(5).map(lambda x: x * 2)
+      fn = original_dataset._trace_variant_creation()
+      variant = fn()
+
+      revived_dataset = dataset_ops._VariantDataset(
+          variant, original_dataset.element_spec)
+      self.assertDatasetProduces(revived_dataset, range(0, 10, 2))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testAsFunctionWithMapInFlatMap(self):
+    with ops.device("CPU"):
+      original_dataset = dataset_ops.Dataset.range(5).flat_map(
+          lambda x: dataset_ops.Dataset.range(5).map(lambda x: x * 2))
+      fn = original_dataset._trace_variant_creation()
+      variant = fn()
+
+      revived_dataset = dataset_ops._VariantDataset(
+          variant, original_dataset.element_spec)
+      self.assertDatasetProduces(revived_dataset, list(original_dataset))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testAsFunctionFromReader(self):
+    with ops.device("CPU"):
+      file_path = os.path.join(self.get_temp_dir(),
+                               "{}.tfrecord.gz".format("tf_record_asset"))
+      with tf_record.TFRecordWriter(file_path, "GZIP") as f:
+        for v in ["a", "aa", "aaa"]:
+          f.write(str(v))
+      original_dataset = readers.TFRecordDataset([file_path],
+                                                 compression_type="GZIP")
+      fn = original_dataset._trace_variant_creation()
+      variant = fn()
+
+      revived_dataset = dataset_ops._VariantDataset(
+          variant, original_dataset.element_spec)
+      self.assertDatasetProduces(revived_dataset, ["a", "aa", "aaa"])
+
+  def _testNumInputs(self, dataset, num_inputs):
+    self.assertLen(dataset._inputs(), num_inputs)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFixedLengthRecordInputs(self):
+    dataset = readers.FixedLengthRecordDataset("", 42)
+    self._testNumInputs(dataset, 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorInputs(self):
+    def gen():
+      yield 42
+
+    dataset = dataset_ops.Dataset.from_generator(gen, dtypes.int32)
+    self._testNumInputs(dataset, 1)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsInputs(self):
+    dataset = dataset_ops.Dataset.from_tensors([42])
+    self._testNumInputs(dataset, 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeInputs(self):
+    dataset = dataset_ops.Dataset.range(10)
+    self._testNumInputs(dataset, 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTextLineInputs(self):
+    dataset = readers.TextLineDataset("")
+    self._testNumInputs(dataset, 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordInputs(self):
+    dataset = readers.TFRecordDataset("")
+    self._testNumInputs(dataset, 1)
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=1, mode=["eager", "graph"]))
+  def testDatasetComplexSourceInputs(self):
+    dataset_fn = dataset_ops.Dataset.from_sparse_tensor_slices(
+        sparse_tensor.SparseTensor(
+            indices=np.array([[0, 0], [1, 0], [2, 0]]),
+            values=np.array([0, 0, 0]),
+            dense_shape=np.array([3, 1])))
+    self.assertEmpty(dataset_fn._inputs())
+
+  def _testUnaryInputs(self, dataset_fn):
+    input_dataset = dataset_ops.Dataset.range(0)
+    self.assertEqual([input_dataset], dataset_fn(input_dataset)._inputs())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBatchInputs(self):
+    self._testUnaryInputs(lambda x: x.batch(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCacheInputs(self):
+    self._testUnaryInputs(lambda x: x.cache())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFilterInputs(self):
+    self._testUnaryInputs(lambda x: x.filter(lambda x: True))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlatMapInputs(self):
+    self._testUnaryInputs(
+        lambda x: x.flat_map(lambda x: dataset_ops.Dataset.range(0)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapInputs(self):
+    self._testUnaryInputs(lambda x: x.map(lambda x: x))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchInputs(self):
+    self._testUnaryInputs(lambda x: x.padded_batch(10, []))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelMapInputs(self):
+    self._testUnaryInputs(lambda x: x.map(lambda x: x, num_parallel_calls=2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRepeatInputs(self):
+    self._testUnaryInputs(lambda x: x.repeat())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleInputs(self):
+    self._testUnaryInputs(lambda x: x.shuffle(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSkipInputs(self):
+    self._testUnaryInputs(lambda x: x.skip(1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTakeInputs(self):
+    self._testUnaryInputs(lambda x: x.take(1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowInputs(self):
+    self._testUnaryInputs(lambda x: x.window(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnaryTransformationInputsApply(self):
+    input_dataset = dataset_ops.Dataset.range(0)
+    dataset = input_dataset.apply(lambda dataset: dataset.cache())
+
+    self.assertEqual([input_dataset], dataset._inputs())
+
+  def _testInputsWithInterleaveFn(self, dataset_fn, interleave_parallelism):
+    input_dataset = dataset_ops.Dataset.range(0)
+    dataset = input_dataset.interleave(
+        lambda x: dataset_ops.Dataset.range(0),
+        cycle_length=2,
+        num_parallel_calls=interleave_parallelism)
+    self.assertEqual([input_dataset], dataset._inputs())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelInterleaveInputs(self):
+    self._testInputsWithInterleaveFn(lambda: dataset_ops.range(0), 2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInterleaveInputs(self):
+    self._testInputsWithInterleaveFn(lambda: dataset_ops.range(0), None)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDebugString(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.map(lambda x: x**2)
+    dataset = dataset.filter(lambda x: x > 10)
+    debug_string = dataset.__debug_string__()
+    for transformation in ["Range", "Map", "Filter"]:
+      self.assertIn(transformation, debug_string)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoWarnings(self):
+    with test.mock.patch.object(warnings, "warn") as mock_log:
+      dataset_ops.Dataset.range(0).interleave(
+          lambda x: dataset_ops.Dataset.range(0), cycle_length=2)
+      self.assertEmpty(mock_log.call_args_list)
+
+  def _testBinaryInputs(self, dataset_fn):
+    input1 = dataset_ops.Dataset.range(0)
+    input2 = dataset_ops.Dataset.range(1)
+    self.assertEqual([input1, input2], dataset_fn(input1, input2)._inputs())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcatenateInputs(self):
+    self._testBinaryInputs(lambda x, y: x.concatenate(y))
+
+  def _testVariadicInputs(self, dataset_fn, input_datasets):
+    self.assertEqual(
+        nest.flatten(input_datasets),
+        dataset_fn(input_datasets)._inputs())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipOneInputs(self):
+    input_datasets = dataset_ops.Dataset.range(0)
+    self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipNestInputs(self):
+    input_datasets = (dataset_ops.Dataset.range(0),
+                      (dataset_ops.Dataset.range(1),
+                       dataset_ops.Dataset.range(2)))
+    self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipTupleInputs(self):
+    input_datasets = (dataset_ops.Dataset.range(0),
+                      dataset_ops.Dataset.range(1))
+    self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFunctions(self):
+    dataset = dataset_ops.Dataset.range(5).map(lambda x: x * 2)
+    self.assertLen(dataset._functions(), 1)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCollectInputs(self):
+    ds1 = dataset_ops.Dataset.range(0)
+    ds2 = ds1.concatenate(ds1)
+    ds3 = dataset_ops.Dataset.zip((ds2, ds1, ds2))
+
+    inputs = []
+    queue = [ds3]
+    while queue:
+      ds = queue[0]
+      queue = queue[1:]
+      queue.extend(ds._inputs())
+      inputs.append(ds)
+
+    self.assertEqual(5, inputs.count(ds1))
+    self.assertEqual(2, inputs.count(ds2))
+    self.assertEqual(1, inputs.count(ds3))
+
+  def _testDatasetSpec(self, tf_value, expected_element_structure):
+    dataset = dataset_ops.Dataset.from_tensors(0).map(lambda _: tf_value)
+    dataset_structure = structure.type_spec_from_value(dataset)
+    self.assertIsInstance(dataset_structure, dataset_ops.DatasetSpec)
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(dataset), expected_element_structure))
+    self.assertEqual([dtypes.variant],
+                     structure.get_flat_tensor_types(dataset_structure))
+    self.assertEqual([tensor_shape.TensorShape([])],
+                     structure.get_flat_tensor_shapes(dataset_structure))
+
+    # Assert that the `Dataset` survives a round-trip via _from_tensor_list()
+    # and _to_tensor_list().
+    round_trip_dataset = dataset_structure._from_tensor_list(
+        dataset_structure._to_tensor_list(dataset))
+
+    value = tf_value
+
+    if isinstance(value, dataset_ops.Dataset):
+      self.assertDatasetsEqual(value, dataset.flat_map(lambda x: x))
+    elif isinstance(value, optional_ops.Optional):
+      self.assertDatasetProduces(
+          round_trip_dataset.map(lambda opt: opt.get_value()),
+          [self.evaluate(value.get_value())],
+          requires_initialization=True)
+    else:
+      self.assertDatasetProduces(
+          round_trip_dataset, [self.evaluate(tf_value)],
+          requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorDatasetSpec(self):
+    self._testDatasetSpec(
+        constant_op.constant(37.0), tensor_spec.TensorSpec([], dtypes.float32))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparseTensorDatasetSpec(self):
+    self._testDatasetSpec(
+        sparse_tensor.SparseTensor(
+            indices=[[0]],
+            values=constant_op.constant([0], dtype=dtypes.int32),
+            dense_shape=[1]), sparse_tensor.SparseTensorSpec([1], dtypes.int32))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestDatasetSpec(self):
+    self._testDatasetSpec(
+        {
+            "a": constant_op.constant(37.0),
+            "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar"))
+        }, {
+            "a":
+                tensor_spec.TensorSpec([], dtypes.float32),
+            "b": (
+                tensor_spec.TensorSpec([1], dtypes.string),
+                tensor_spec.TensorSpec([], dtypes.string),
+            )
+        })
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetDatasetSpec(self):
+    self._testDatasetSpec(
+        dataset_ops.Dataset.from_tensor_slices(
+            constant_op.constant([1, 2, 3])),
+        dataset_ops.DatasetSpec(tensor_spec.TensorSpec([], dtypes.int32)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionalDatasetSpec(self):
+    self._testDatasetSpec(
+        optional_ops.Optional.from_value(37.0),
+        optional_ops.OptionalSpec(tensor_spec.TensorSpec([], dtypes.float32)))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSameGraphError(self):
+    dataset = dataset_ops.Dataset.range(10)
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(ValueError, "must be from the same graph"):
+        dataset = dataset.batch(2)
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[1], mode=["graph"]))
+  def testSameGraphErrorOneShot(self):
+    dataset = dataset_ops.Dataset.range(10)
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(ValueError,
+                                  "Make sure that the dataset is created in "
+                                  "the same graph as the iterator"):
+        _ = dataset_ops.make_one_shot_iterator(dataset)
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[1], mode=["graph"]))
+  def testSameGraphErrorInitializable(self):
+    dataset = dataset_ops.Dataset.range(10)
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(ValueError,
+                                  "Make sure that the dataset is created in "
+                                  "the same graph as the iterator"):
+        _ = dataset_ops.make_initializable_iterator(dataset)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(execution_mode=[context.ASYNC, context.SYNC])))
+  def testEagerIteration(self, execution_mode):
+    with context.execution_mode(execution_mode):
+      val = 0
+      dataset = dataset_ops.Dataset.range(10)
+      for foo in dataset:
+        self.assertEqual(val, foo.numpy())
+        val += 1
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetAsFunctionArgument(self):
+
+    @def_function.function
+    def _uses_dataset(d):
+      accumulator = array_ops.zeros([], dtype=dtypes.int64)
+      for value in d:
+        accumulator += value
+      return accumulator
+
+    with ops.device("CPU"):
+      first_dataset = dataset_ops.Dataset.range(10)
+      self.assertEqual(45, self.evaluate(_uses_dataset(first_dataset)))
+      second_dataset = dataset_ops.Dataset.range(11)
+      self.assertEqual(55, self.evaluate(_uses_dataset(second_dataset)))
+      first_concrete = _uses_dataset.get_concrete_function(first_dataset)
+      # The dataset should not be a captured input
+      self.assertEmpty(first_concrete.graph.captures)
+      # The two datasets have the same structure and so should re-use a trace.
+      self.assertIs(first_concrete,
+                    _uses_dataset.get_concrete_function(second_dataset))
+      # With a different structure we should use a different trace.
+      self.assertIsNot(
+          first_concrete,
+          _uses_dataset.get_concrete_function(
+              dataset_ops.Dataset.zip((first_dataset, second_dataset))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLimitedRetracing(self):
+    trace_count = [0]
+
+    @def_function.function
+    def f(ds):
+      trace_count[0] += 1
+      counter = np.int64(0)
+      for elem in ds:
+        counter += elem
+      return counter
+
+    dataset = dataset_ops.Dataset.range(5)
+    dataset2 = dataset_ops.Dataset.range(10)
+
+    for _ in range(10):
+      self.assertEqual(self.evaluate(f(dataset)), 10)
+      self.assertEqual(self.evaluate(f(dataset2)), 45)
+      self.assertEqual(trace_count[0], 1)
+
+  # pylint: disable=g-long-lambda,unnecessary-lambda
+  @combinations.generate(test_base.default_test_combinations())
+  def testLegacyStructureAPI(self):
+    components = (np.array([1, 2, 3], dtype=np.int64), (np.array([4., 5.]),
+                                                        np.array([6., 7.])),
+                  np.array([8, 9, 10], dtype=np.int64))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([3], ([2], [2]), [3]),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.shuffle(10, 10)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([3], ([2], [2]), [3]),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.repeat(-1)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([3], ([2], [2]), [3]),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.filter(lambda x, y, z: True)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([3], ([2], [2]), [3]),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.take(5)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([3], ([2], [2]), [3]),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.map(lambda x, y, z: ((x, z), (y[0], y[1])))
+    self.assertEqual(
+        ((dtypes.int64, dtypes.int64), (dtypes.float64, dtypes.float64)),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual((([3], [3]), ([2], [2])),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.flat_map(lambda x, y: dataset_ops.Dataset.from_tensors(
+        ((x[0], x[1]), (y[0], y[1]))))
+    self.assertEqual(
+        ((dtypes.int64, dtypes.int64), (dtypes.float64, dtypes.float64)),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual((([3], [3]), ([2], [2])),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.batch(32)
+    self.assertEqual(
+        ((dtypes.int64, dtypes.int64), (dtypes.float64, dtypes.float64)),
+        dataset_ops.get_legacy_output_types(dataset))
+    dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset)
+    self.assertEqual(
+        (([None, 3], [None, 3]), ([None, 2], [None, 2])),
+        nest.pack_sequence_as(
+            dataset_output_shapes,
+            [s.as_list() for s in nest.flatten(dataset_output_shapes)]))
+
+    # Define a separate set of components with matching leading
+    # dimension for the from-slices constructor.
+    components_for_slices = (np.array([1, 2, 3],
+                                      dtype=np.int64), (np.array([4., 5., 6.]),
+                                                        np.array([7., 8., 9.])),
+                             np.array([10, 11, 12], dtype=np.int64))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components_for_slices)
+    self.assertEqual(
+        (dtypes.int64, (dtypes.float64, dtypes.float64), dtypes.int64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual(([], ([], []), []),
+                     dataset_ops.get_legacy_output_shapes(dataset))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoneComponent(self):
+    dataset = dataset_ops.Dataset.from_tensors((42, None))
+    if context.executing_eagerly():
+      self.assertDatasetProduces(dataset, expected_output=[(42, None)])
+    else:
+      iterator = dataset_ops.make_one_shot_iterator(dataset)
+      next_first, next_second = iterator.get_next()
+      self.assertIsNone(next_second)
+      with self.cached_session() as sess:
+        self.assertEqual(sess.run(next_first), 42)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoneComponentInFunction(self):
+
+    @def_function.function
+    def fn(ds):
+      total = 0
+      it = iter(ds)
+      for elem in it:
+        x, _ = elem
+        total += x
+      return total
+
+    dataset = dataset_ops.Dataset.range(
+        10, output_type=dtypes.int32).map(lambda x: (x, None))
+    self.assertEqual(self.evaluate(fn(dataset)), 45)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIncorrectPythonStructure(self):
+    # Tests that an exception is raised (as opposed to a segfault) when the
+    # Python structure assigned to a dataset is incorrect.
+    dataset = dataset_ops.Dataset.range(10)
+    spec = tensor_spec.TensorSpec([], dtypes.int64)
+    new_structure = (spec, spec)
+    dataset = dataset_ops._RestructuredDataset(dataset, new_structure)
+    dataset = dataset.map(lambda x, y: y)
+
+    with self.assertRaisesOpError(""):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNamedTupleStructure(self):
+    Foo = collections.namedtuple("Foo", ["a", "b"])
+    x = Foo(a=3, b="test")
+    dataset = dataset_ops.Dataset.from_tensors(x)
+    dataset = dataset_ops.Dataset.from_tensor_slices([dataset, dataset])
+    self.assertEqual(
+        str(dataset.element_spec),
+        "DatasetSpec(Foo(a=TensorSpec(shape=(), dtype=tf.int32, name=None), "
+        "b=TensorSpec(shape=(), dtype=tf.string, name=None)), TensorShape([]))")
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testIterationError(self):
+
+    @def_function.function(autograph=False)
+    def fn(ds):
+      for _ in ds:
+        pass
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaises(ValueError):
+      self.evaluate(fn(dataset))
+
+
+class DebugDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(DebugDatasetTest, self).setUp()
+    debug_mode.toggle_debug_mode(True)
+
+  def tearDown(self):
+    debug_mode.toggle_debug_mode(False)
+    super(DebugDatasetTest, self).tearDown()
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDebugModeEagerExecution(self):
+    counter = []
+    ds = dataset_ops.Dataset.range(10)
+
+    def map_fn(x):
+      counter.append(1)
+      return x
+
+    ds = ds.map(map_fn)
+    self.assertDatasetProduces(ds, list(range(10)))
+
+    # The body of `map_fn` will be executed 11 times since the implementation
+    # traces the function to figure out what the types and shapes of its
+    # outputs are.
+    self.assertLen(counter, 11)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDebugModeGenerator(self):
+    def gen():
+      yield from range(10)
+
+    ds = dataset_ops.Dataset.from_generator(
+        gen,
+        output_signature=tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64))
+    self.assertDatasetProduces(ds, list(range(10)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDebugModeGeneratorTwoComponents(self):
+    data = [(n, n+1) for n in range(10)]
+
+    def gen():
+      yield from data
+
+    ds = dataset_ops.Dataset.from_generator(
+        gen,
+        output_signature=(tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64),
+                          tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64)))
+    self.assertDatasetProduces(ds, data)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDebugModeSequentialExecution(self):
+    ds = dataset_ops.Dataset.range(10)
+    ds = ds.apply(
+        testing.assert_next(["Interleave", "Map", "Batch", "FiniteTake"]))
+    ds = ds.interleave(
+        dataset_ops.Dataset.from_tensors,
+        cycle_length=10,
+        num_parallel_calls=10)
+    ds = ds.map(lambda x: x * x, num_parallel_calls=10)
+    ds = ds.batch(batch_size=5, num_parallel_calls=2)
+    ds = ds.prefetch(buffer_size=2)
+    ds = ds.take(2)
+    self.assertDatasetProduces(ds, [[0, 1, 4, 9, 16], [25, 36, 49, 64, 81]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/enumerate_test.py

@@ -0,0 +1,75 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.enumerate()`."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+
+
+class EnumerateTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEnumerate(self):
+    components = (["a", "b"], [1, 2], [37.0, 38])
+    start = constant_op.constant(20, dtype=dtypes.int64)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).enumerate(
+        start)
+
+    self.assertEqual(dtypes.int64,
+                     dataset_ops.get_legacy_output_types(dataset)[0])
+    dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset)
+    self.assertEqual((), dataset_output_shapes[0])
+    self.assertEqual([tensor_shape.TensorShape([])] * 3,
+                     [shape for shape in dataset_output_shapes[1]])
+
+    self.assertDatasetProduces(dataset, [(20, (b"a", 1, 37.0)),
+                                         (21, (b"b", 2, 38.0))])
+
+
+class EnumerateCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                              parameterized.TestCase):
+
+  def _build_enumerate_dataset(self, start, stop, options=None):
+    dataset = dataset_ops.Dataset.range(start, stop).enumerate()
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    start = 2
+    stop = 10
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self, lambda: self._build_enumerate_dataset(
+            start=start, stop=stop, options=options), stop - start)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/filter_test.py

@@ -0,0 +1,298 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.filter()`."""
+
+from typing import Callable
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+
+  def filter_fn(dataset, predicate):
+    return dataset.filter(predicate)
+
+  def legacy_filter_fn(dataset, predicate):
+    return dataset.filter_with_legacy_function(predicate)
+
+  filter_combinations = combinations.combine(
+      tf_api_version=[1, 2],
+      mode=["eager", "graph"],
+      apply_filter=combinations.NamedObject("filter_fn", filter_fn))
+
+  legacy_filter_combinations = combinations.combine(
+      tf_api_version=1,
+      mode=["eager", "graph"],
+      apply_filter=combinations.NamedObject("legacy_filter_fn",
+                                            legacy_filter_fn))
+
+  return filter_combinations + legacy_filter_combinations
+
+
+class FilterTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(_test_combinations())
+  def testFilterDataset(self, apply_filter):
+    components = (np.arange(7, dtype=np.int64),
+                  np.array([[1, 2, 3]], dtype=np.int64) *
+                  np.arange(7, dtype=np.int64)[:, np.newaxis],
+                  np.array(37.0, dtype=np.float64) * np.arange(7))
+
+    def _map_fn(x, y, z):
+      return math_ops.square(x), math_ops.square(y), math_ops.square(z)
+
+    def do_test(count, modulus):  # pylint: disable=missing-docstring
+      dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+          _map_fn).repeat(count)
+      # pylint: disable=g-long-lambda
+      dataset = apply_filter(
+          dataset,
+          lambda x, _y, _z: math_ops.equal(math_ops.mod(x, modulus), 0))
+      # pylint: enable=g-long-lambda
+      self.assertEqual(
+          [c.shape[1:] for c in components],
+          [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+      get_next = self.getNext(dataset)
+      for _ in range(count):
+        for i in [x for x in range(7) if x**2 % modulus == 0]:
+          result = self.evaluate(get_next())
+          for component, result_component in zip(components, result):
+            self.assertAllEqual(component[i]**2, result_component)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+    do_test(14, 2)
+    do_test(4, 18)
+
+    # Test an empty dataset.
+    do_test(0, 1)
+
+  @combinations.generate(_test_combinations())
+  def testFilterRange(self, apply_filter):
+    dataset = dataset_ops.Dataset.range(4)
+    dataset = apply_filter(dataset,
+                           lambda x: math_ops.not_equal(math_ops.mod(x, 3), 2))
+    self.assertDatasetProduces(dataset, expected_output=[0, 1, 3])
+
+  @combinations.generate(_test_combinations())
+  def testFilterDict(self, apply_filter):
+    dataset = dataset_ops.Dataset.range(10).map(
+        lambda x: {"foo": x * 2, "bar": x**2})
+    dataset = apply_filter(dataset, lambda d: math_ops.equal(d["bar"] % 2, 0))
+    dataset = dataset.map(lambda d: d["foo"] + d["bar"])
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[(i * 2 + i**2) for i in range(10) if not (i**2) % 2])
+
+  @combinations.generate(_test_combinations())
+  def testUseStepContainerInFilter(self, apply_filter):
+    input_data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
+
+    # Define a predicate that returns true for the first element of
+    # the sequence and not the second, and uses `tf.map_fn()`.
+    def _predicate(xs):
+      squared_xs = map_fn.map_fn(lambda x: x * x, xs)
+      summed = math_ops.reduce_sum(squared_xs)
+      return math_ops.equal(summed, 1 + 4 + 9)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices([[1, 2, 3], [4, 5, 6]])
+    dataset = apply_filter(dataset, _predicate)
+    self.assertDatasetProduces(dataset, expected_output=[input_data[0]])
+
+  @combinations.generate(_test_combinations())
+  def testSparse(self, apply_filter):
+
+    def _map_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1])),
+          dense_shape=np.array([1, 1])), i
+
+    def _filter_fn(_, i):
+      return math_ops.equal(i % 2, 0)
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn)
+    dataset = apply_filter(dataset, _filter_fn)
+    dataset = dataset.map(lambda x, i: x)
+    self.assertDatasetProduces(
+        dataset, expected_output=[_map_fn(i * 2)[0] for i in range(5)])
+
+  @combinations.generate(_test_combinations())
+  def testShortCircuit(self, apply_filter):
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.range(10),
+         dataset_ops.Dataset.from_tensors(True).repeat(None)))
+    dataset = apply_filter(dataset, lambda x, y: y)
+    self.assertDatasetProduces(
+        dataset, expected_output=[(i, True) for i in range(10)])
+
+  @combinations.generate(_test_combinations())
+  def testParallelFilters(self, apply_filter):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = apply_filter(dataset, lambda x: math_ops.equal(x % 2, 0))
+    next_elements = [self.getNext(dataset) for _ in range(10)]
+    self.assertEqual([0 for _ in range(10)],
+                     self.evaluate(
+                         [next_element() for next_element in next_elements]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).filter(
+        lambda x: True, name="filter")
+    self.assertDatasetProduces(dataset, [42])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPredicateFailWithErrorContext(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).filter(
+        lambda x: (x // 0) > 0, name="filter")
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        r".*Error in user-defined function passed to .* transformation with "
+        r"iterator: Iterator::Root::.*"):
+      self.evaluate(get_next())
+
+
+class FilterCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                           parameterized.TestCase):
+
+  def _build_filter_range_dataset(self, div, options=None):
+    dataset = dataset_ops.Dataset.range(100).filter(
+        lambda x: math_ops.not_equal(math_ops.mod(x, div), 2))
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    div = 3
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    num_outputs = sum(x % 3 != 2 for x in range(100))
+    verify_fn(self, lambda: self._build_filter_range_dataset(div, options),
+              num_outputs)
+
+  def _build_filter_dict_dataset(self):
+    return dataset_ops.Dataset.range(10).map(lambda x: {
+        "foo": x * 2,
+        "bar": x**2
+    }).filter(lambda d: math_ops.equal(d["bar"] % 2, 0)).map(
+        lambda d: d["foo"] + d["bar"])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testDict(self, verify_fn):
+    num_outputs = sum((x**2) % 2 == 0 for x in range(10))
+    verify_fn(self, self._build_filter_dict_dataset, num_outputs)
+
+  def _build_sparse_filter_dataset(self):
+
+    def _map_fn(i):
+      return sparse_tensor.SparseTensor(
+          indices=[[0, 0]], values=(i * [1]), dense_shape=[1, 1]), i
+
+    def _filter_fn(_, i):
+      return math_ops.equal(i % 2, 0)
+
+    return dataset_ops.Dataset.range(10).map(_map_fn).filter(_filter_fn).map(
+        lambda x, i: x)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testSparse(self, verify_fn):
+    verify_fn(self, self._build_sparse_filter_dataset, num_outputs=5)
+
+
+class FilterGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleFilter(self):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = global_shuffle_op._global_shuffle(dataset)
+    dataset = dataset.filter(lambda x: math_ops.equal(x % 2, 0))
+    self.assertDatasetProduces(
+        dataset,
+        list(range(0, 100, 2)),
+        requires_initialization=True,
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFilterShuffle(self):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset.filter(lambda x: math_ops.equal(x % 2, 0))
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        "`global_shuffle` requires all upstream transformations be compatible "
+        "with random access."):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class FilterGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testShuffleFilter(
+      self,
+      verify_fn: Callable[..., None],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(10)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      dataset = dataset.filter(lambda x: math_ops.equal(x % 2, 0))
+      if symbolic_checkpoint:
+        options = options_lib.Options()
+        options.experimental_symbolic_checkpoint = symbolic_checkpoint
+        dataset = dataset.with_options(options)
+      return dataset
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=5,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/fingerprint_test.py

@@ -0,0 +1,111 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `dataset.fingerprint()`."""
+
+from absl.testing import parameterized
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+def dataset_fn_test_cases():
+  cases = [
+      ("range1", lambda: dataset_ops.Dataset.range(10)),
+      (
+          "flat_map1",
+          lambda: dataset_ops.Dataset.range(10).flat_map(
+              dataset_ops.Dataset.range
+          ),
+      ),
+      ("tfrecord1", lambda: readers.TFRecordDataset(["f1.txt", "f2.txt"])),
+      (
+          "tfrecord2",
+          lambda: readers.TFRecordDataset(["f1.txt", "f2.txt"]).repeat(2),
+      ),
+  ]
+
+  named_cases = []
+  for case in cases:
+    name, dataset_fn = case
+    named_cases.append(combinations.NamedObject(name=name, obj=dataset_fn))
+
+  return combinations.combine(dataset_fn=named_cases)
+
+
+def dataset_pair_fn_test_cases():
+  dataset = dataset_ops.Dataset
+
+  cases = [
+      ("range1", lambda: (dataset.range(10), dataset.range(11))),
+      (
+          "flat_map1",
+          lambda: (
+              dataset.range(10).flat_map(dataset.range),
+              dataset.range(10).flat_map(lambda x: dataset.range(x + 1)),
+          ),
+      ),
+      (
+          "tfrecord1",
+          lambda: (
+              readers.TFRecordDataset(["f1.txt", "f2.txt"]),
+              readers.TFRecordDataset(["f1.txt", "f3.txt"]),
+          ),
+      ),
+      (
+          "tfrecord2",
+          lambda: (
+              readers.TFRecordDataset(["f1.txt", "f2.txt"]).repeat(2),
+              readers.TFRecordDataset(["f1.txt", "f3.txt"]).repeat(2),
+          ),
+      ),
+  ]
+
+  named_cases = []
+  for case in cases:
+    name, dataset_pair_fn = case
+    named_cases.append(combinations.NamedObject(name=name, obj=dataset_pair_fn))
+
+  return combinations.combine(dataset_pair_fn=named_cases)
+
+
+class FingerprintTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(), dataset_fn_test_cases()
+      )
+  )
+  def testSameDatasetSameFingerprint(self, dataset_fn):
+    fingerprint1 = self.evaluate(dataset_fn().fingerprint())
+    fingerprint2 = self.evaluate(dataset_fn().fingerprint())
+    self.assertEqual(fingerprint1, fingerprint2)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          dataset_pair_fn_test_cases(),
+      )
+  )
+  def testDifferentDatasetDifferentFingerprint(self, dataset_pair_fn):
+    lhs, rhs = dataset_pair_fn()
+    lhs_fingerprint = self.evaluate(lhs.fingerprint())
+    rhs_fingerprint = self.evaluate(rhs.fingerprint())
+    self.assertNotEqual(lhs_fingerprint, rhs_fingerprint)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/fixed_length_record_dataset_test.py

@@ -0,0 +1,249 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.FixedLengthRecordDataset`."""
+import gzip
+import os
+import pathlib
+import zlib
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class FixedLengthRecordDatasetTestBase(test_base.DatasetTestBase):
+  """Base class for setting up and testing FixedLengthRecordDataset."""
+
+  def setUp(self):
+    super(FixedLengthRecordDatasetTestBase, self).setUp()
+    self._num_files = 2
+    self._num_records = 7
+    self._header_bytes = 5
+    self._record_bytes = 3
+    self._footer_bytes = 2
+
+  def _record(self, f, r):
+    return compat.as_bytes(str(f * 2 + r) * self._record_bytes)
+
+  def _createFiles(self, compression_type=None):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "fixed_length_record.%d.txt" % i)
+      filenames.append(fn)
+
+      contents = []
+      contents.append(b"H" * self._header_bytes)
+      for j in range(self._num_records):
+        contents.append(self._record(i, j))
+      contents.append(b"F" * self._footer_bytes)
+      contents = b"".join(contents)
+
+      if not compression_type:
+        with open(fn, "wb") as f:
+          f.write(contents)
+      elif compression_type == "GZIP":
+        with gzip.GzipFile(fn, "wb") as f:
+          f.write(contents)
+      elif compression_type == "ZLIB":
+        contents = zlib.compress(contents)
+        with open(fn, "wb") as f:
+          f.write(contents)
+      else:
+        raise ValueError("Unsupported compression_type", compression_type)
+
+    return filenames
+
+
+class FixedLengthRecordDatasetTest(FixedLengthRecordDatasetTestBase,
+                                   parameterized.TestCase):
+
+  def _test(self, compression_type=None):
+    test_filenames = self._createFiles(compression_type=compression_type)
+
+    def dataset_fn(filenames, num_epochs, batch_size=None):
+      repeat_dataset = readers.FixedLengthRecordDataset(
+          filenames,
+          self._record_bytes,
+          self._header_bytes,
+          self._footer_bytes,
+          compression_type=compression_type).repeat(num_epochs)
+      if batch_size:
+        return repeat_dataset.batch(batch_size)
+      return repeat_dataset
+
+    # Basic test: read from file 0.
+    self.assertDatasetProduces(
+        dataset_fn([test_filenames[0]], 1),
+        expected_output=[
+            self._record(0, i) for i in range(self._num_records)
+        ])
+
+    # Basic test: read from file 1.
+    self.assertDatasetProduces(
+        dataset_fn([test_filenames[1]], 1),
+        expected_output=[
+            self._record(1, i) for i in range(self._num_records)
+        ])
+
+    # Basic test: read from both files.
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(
+        dataset_fn(test_filenames, 1), expected_output=expected_output)
+
+    # Test repeated iteration through both files.
+    get_next = self.getNext(dataset_fn(test_filenames, 10))
+    for _ in range(10):
+      for j in range(self._num_files):
+        for i in range(self._num_records):
+          self.assertEqual(self._record(j, i), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Test batched and repeated iteration through both files.
+    get_next = self.getNext(dataset_fn(test_filenames, 10, self._num_records))
+    for _ in range(10):
+      for j in range(self._num_files):
+        self.assertAllEqual(
+            [self._record(j, i) for i in range(self._num_records)],
+            self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoCompression(self):
+    self._test()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGzipCompression(self):
+    self._test(compression_type="GZIP")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZlibCompression(self):
+    self._test(compression_type="ZLIB")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBuffering(self):
+    test_filenames = self._createFiles()
+    dataset = readers.FixedLengthRecordDataset(
+        test_filenames,
+        self._record_bytes,
+        self._header_bytes,
+        self._footer_bytes,
+        buffer_size=10)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelRead(self):
+    test_filenames = self._createFiles()
+    dataset = readers.FixedLengthRecordDataset(
+        test_filenames,
+        self._record_bytes,
+        self._header_bytes,
+        self._footer_bytes,
+        buffer_size=10,
+        num_parallel_reads=4)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output,
+                               assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWrongSize(self):
+    test_filenames = self._createFiles()
+    dataset = readers.FixedLengthRecordDataset(
+        test_filenames,
+        self._record_bytes + 1,  # Incorrect record length.
+        self._header_bytes,
+        self._footer_bytes,
+        buffer_size=10)
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(
+            errors.InvalidArgumentError,
+            r"Excluding the header \(5 bytes\) and footer \(2 bytes\), input "
+            r"file \".*fixed_length_record.0.txt\" has body length 21 bytes, "
+            r"which is not an exact multiple of the record length \(4 bytes\).")
+        )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPathlib(self):
+    test_filenames = self._createFiles()
+    test_filenames = [pathlib.Path(f) for f in test_filenames]
+    dataset = readers.FixedLengthRecordDataset(
+        test_filenames,
+        self._record_bytes,
+        self._header_bytes,
+        self._footer_bytes,
+        buffer_size=10,
+        num_parallel_reads=4)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output,
+                               assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    test_filenames = self._createFiles()
+    dataset = readers.FixedLengthRecordDataset(
+        test_filenames,
+        self._record_bytes,
+        self._header_bytes,
+        self._footer_bytes,
+        name="fixed_length_record_dataset")
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+
+class FixedLengthRecordDatasetCheckpointTest(
+    FixedLengthRecordDatasetTestBase, checkpoint_test_base.CheckpointTestBase,
+    parameterized.TestCase):
+
+  def _build_dataset(self, num_epochs, compression_type=None):
+    filenames = self._createFiles()
+    return readers.FixedLengthRecordDataset(
+        filenames, self._record_bytes, self._header_bytes,
+        self._footer_bytes).repeat(num_epochs)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    num_epochs = 5
+    num_outputs = num_epochs * self._num_files * self._num_records
+    verify_fn(self, lambda: self._build_dataset(num_epochs), num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/flat_map_test.py

@@ -0,0 +1,564 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.flat_map()`."""
+import random
+from typing import Callable, Optional
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.framework import dataset_options_pb2
+from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import function
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops.ragged import ragged_conversion_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class FlatMapTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  # pylint: disable=g-long-lambda
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlatMapDataset(self):
+    repeats = [1, 2, 3, 4, 5, 0, 1]
+    components = np.array(repeats, dtype=np.int64)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).flat_map(
+        lambda x: dataset_ops.Dataset.from_tensors([x]).repeat(x)
+    )
+    expected_output = []
+    for i in repeats:
+      expected_output.extend([[i]] * i)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedFlatMapDataset(self):
+    repeats = [[1, 2], [3, 4], [5, 0], [1, 7]]
+    components = np.array(repeats, dtype=np.int64)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).flat_map(
+        lambda x: dataset_ops.Dataset.from_tensor_slices(x).flat_map(
+            lambda y: dataset_ops.Dataset.from_tensors(y).repeat(y)
+        )
+    )
+    expected_output = []
+    for row in repeats:
+      for i in row:
+        expected_output.extend([i] * i)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSharedResourceNestedFlatMapDataset(self):
+    repeats = [[1, 2], [3, 4], [5, 0], [1, 7]]
+    components = np.array(repeats, dtype=np.int64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_tensor_slices(components).flat_map(
+            lambda x: dataset_ops.Dataset.from_tensor_slices(x).flat_map(
+                lambda y: dataset_ops.Dataset.from_tensors(y).repeat(y)
+            )
+        ),
+        shared_name="shared_flat_map_iterator",
+    )
+    init_op = iterator.initializer
+    get_next = iterator.get_next()
+
+    # Create two concurrent sessions that share the same iterator
+    # resource on the same server, and verify that a random
+    # interleaving of `Session.run(get_next)` calls on the two
+    # sessions yields the expected result.
+    server = server_lib.Server.create_local_server()
+    with session.Session(server.target) as sess1:
+      with session.Session(server.target) as sess2:
+        for _ in range(3):
+          sess = random.choice([sess1, sess2])
+          sess.run(init_op)
+          for row in repeats:
+            for i in row:
+              for _ in range(i):
+                sess = random.choice([sess1, sess2])
+                self.assertEqual(i, sess.run(get_next))
+
+        with self.assertRaises(errors.OutOfRangeError):
+          sess = random.choice([sess1, sess2])
+          sess.run(get_next)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapDict(self):
+    dataset = (
+        dataset_ops.Dataset.range(10)
+        .map(lambda x: {"foo": x * 2, "bar": x**2})
+        .flat_map(
+            lambda d: dataset_ops.Dataset.from_tensors(d["foo"]).repeat(
+                d["bar"]
+            )
+        )
+    )
+    get_next = self.getNext(dataset)
+    for i in range(10):
+      for _ in range(i**2):
+        self.assertEqual(i * 2, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparse(self):
+    def _map_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]
+      )
+
+    def _flat_map_fn(x):
+      return dataset_ops.Dataset.from_tensor_slices(
+          sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)
+      )
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn)
+    expected_output = []
+    for i in range(10):
+      for j in range(2):
+        expected_output.append([i, 0] if j % 2 == 0 else [0, -i])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorArray(self):
+    def _map_fn(i):
+      i = math_ops.cast(i, dtypes.int32)
+      return tensor_array_ops.TensorArray(
+          dtype=dtypes.int32, element_shape=(), size=i
+      ).unstack(math_ops.range(i))
+
+    def _flat_map_fn(x):
+      self.assertIsInstance(x, tensor_array_ops.TensorArray)
+      return dataset_ops.Dataset.from_tensor_slices(x.stack())
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn)
+
+    expected_output = []
+    for i in range(10):
+      for j in range(i):
+        expected_output.append(j)
+
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRagged(self):
+    def _map_fn(i):
+      return ragged_tensor.RaggedTensor.from_tensor(i * [[1], [-1]])
+
+    def _flat_map_fn(x):
+      return dataset_ops.Dataset.from_tensor_slices(
+          ragged_conversion_ops.to_tensor(x)
+      )
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn)
+    expected_output = []
+    for i in range(10):
+      expected_output.append([i])
+      expected_output.append([-i])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    def fn(x):
+      return dataset_ops.Dataset.from_tensors(x)
+
+    dataset = dataset_ops.Dataset.from_tensors(42).flat_map(fn, name="flat_map")
+    self.assertDatasetProduces(dataset, [42])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCardinality(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+    options = dataset_options_pb2.CardinalityOptions(
+        compute_level="CARDINALITY_COMPUTE_MODERATE")
+    cardinality = dataset_ops.gen_dataset_ops.dataset_cardinality(
+        dataset._variant_tensor, options.SerializeToString())
+    self.assertEqual(self.evaluate(cardinality), 9)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteCardinality(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    dataset = dataset.flat_map(lambda _: dataset_ops.Dataset.range(1).repeat())
+    options = dataset_options_pb2.CardinalityOptions(
+        compute_level="CARDINALITY_COMPUTE_MODERATE")
+    cardinality = dataset_ops.gen_dataset_ops.dataset_cardinality(
+        dataset._variant_tensor, options.SerializeToString())
+    self.assertEqual(self.evaluate(cardinality), dataset_ops.INFINITE)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownCardinality(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    dataset = dataset.flat_map(
+        lambda _: dataset_ops.Dataset.range(10).filter(lambda x: x % 2 == 1))
+    options = dataset_options_pb2.CardinalityOptions(
+        compute_level="CARDINALITY_COMPUTE_MODERATE")
+    cardinality = dataset_ops.gen_dataset_ops.dataset_cardinality(
+        dataset._variant_tensor, options.SerializeToString())
+    self.assertEqual(self.evaluate(cardinality), dataset_ops.UNKNOWN)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyCardinality(self):
+    dataset = dataset_ops.Dataset.range(0)
+    dataset = dataset.flat_map(dataset_ops.Dataset.range)
+    options = dataset_options_pb2.CardinalityOptions(
+        compute_level="CARDINALITY_COMPUTE_MODERATE")
+    cardinality = dataset_ops.gen_dataset_ops.dataset_cardinality(
+        dataset._variant_tensor, options.SerializeToString())
+    self.assertEqual(self.evaluate(cardinality), 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCardinalityLowEffort(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+    options = dataset_options_pb2.CardinalityOptions(
+        compute_level="CARDINALITY_COMPUTE_LOW")
+    cardinality = dataset_ops.gen_dataset_ops.dataset_cardinality(
+        dataset._variant_tensor, options.SerializeToString())
+    self.assertEqual(self.evaluate(cardinality), dataset_ops.UNKNOWN)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapFuncFailWithErrorContext(self):
+
+    def fn(x):
+      return dataset_ops.Dataset.from_tensors(x // 0)
+
+    dataset = dataset_ops.Dataset.from_tensors(42).flat_map(fn, name="flat_map")
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        r".*Error in user-defined function passed to .* transformation with "
+        r"iterator: Iterator::Root::.*"):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testSymbolicCheckpointSize(self):
+    examples_per_flat_map = 100
+    example_len = 10_000
+
+    def flat_map_fn(_):
+      data = []
+      for _ in range(examples_per_flat_map):
+        data.append(
+            stateless_random_ops.stateless_random_uniform(
+                [example_len], seed=(42, 42)
+            )
+        )
+      return dataset_ops.Dataset.from_tensor_slices(data)
+
+    ds = dataset_ops.Dataset.range(10)
+    # Inputs to flat_map are >1MB
+    ds = ds.map(
+        lambda x: stateless_random_ops.stateless_random_uniform(
+            [1_000_000], seed=(42, 42)
+        )
+    )
+    ds = ds.flat_map(flat_map_fn)
+
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = True
+    ds = ds.with_options(options)
+
+    it = ds.as_numpy_iterator()
+    for _ in range(30):
+      next(it)
+
+    ckpt = it._save()
+    # Make sure the checkpoint is smaller than the element sizes, i.e. no
+    # elements are being stored in the checkpoint.
+    self.assertLess(len(ckpt.numpy()), 10_000)
+
+
+class FlatMapCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True]),
+      )
+  )
+  def test(self, verify_fn, symbolic_checkpoint):
+    # Complicated way of saying range(start, start+25).
+    def build_ds(start):
+      def map_fn(x):
+        return dataset_ops.Dataset.range(x, x + 5)
+
+      dataset = dataset_ops.Dataset.range(start, start + 5 * 5, 5)
+      dataset = dataset.flat_map(map_fn)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(self, lambda: build_ds(0), num_outputs=25)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def testNested(self, verify_fn):
+    def build_ds():
+      inner_ds = dataset_ops.Dataset.from_tensor_slices(range(42))
+      ds = dataset_ops.Dataset.from_tensors(inner_ds)
+      return ds.flat_map(lambda x: x)
+
+    verify_fn(self, build_ds, num_outputs=42)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def testMapThenFlatMap(self, verify_fn):
+    def build_ds():
+      def flat_map_fn(_):
+        def map_fn(y):
+          return 10 * math_ops.cast(y, dtypes.int32)
+
+        return dataset_ops.Dataset.range(100).map(map_fn)
+
+      return dataset_ops.Dataset.range(5).flat_map(flat_map_fn)
+
+    verify_fn(self, build_ds, num_outputs=500)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def testCaptureDefunInMapFn(self, verify_fn):
+    def build_ds():
+      def map_fn(x):
+        @function.Defun(dtypes.int64)
+        def defun_fn(x):
+          return constant_op.constant(1000) + math_ops.cast(x, dtypes.int32)
+
+        return dataset_ops.Dataset.from_tensor_slices([defun_fn(x)])
+
+      return dataset_ops.Dataset.range(100).flat_map(map_fn)
+
+    verify_fn(self, build_ds, num_outputs=100)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDisallowVariableCapture(self):
+    def build_ds():
+      test_var = variable_scope.get_variable(
+          name="test_var", shape=(), use_resource=True
+      )
+      return dataset_ops.Dataset.range(5).flat_map(
+          lambda _: dataset_ops.Dataset.from_tensor_slices([test_var])
+      )
+
+    self.verify_error_on_save(build_ds, 5, errors.FailedPreconditionError)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDisallowCapturingStatefulOps(self):
+    def build_ds():
+      def flat_map_fn(_):
+        def map_fn(x):
+          return random_ops.random_uniform(
+              (), 0, 10, dtype=dtypes.int32
+          ) * math_ops.cast(x, dtypes.int32)
+
+        return dataset_ops.Dataset.range(100).map(map_fn)
+
+      return dataset_ops.Dataset.range(5).flat_map(flat_map_fn)
+
+    self.verify_error_on_save(build_ds, 500, errors.FailedPreconditionError)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def testSparse(self, verify_fn):
+    def _map_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2]
+      )
+
+    def _flat_map_fn(x):
+      return dataset_ops.Dataset.from_tensor_slices(
+          sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values)
+      )
+
+    def _build_ds():
+      return dataset_ops.Dataset.range(10).map(_map_fn).flat_map(_flat_map_fn)
+
+    verify_fn(self, _build_ds, num_outputs=20)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[True],
+                               num_skips=[3, 4]),
+      )
+  )
+  def testWithSkip(self, verify_fn, symbolic_checkpoint, num_skips):
+    """Test `.flat_map().skip()` checkpointing behavior.
+
+    `SkipInternal` and `GetNextInternal` are separate functions
+    but with slightly different implementations.
+    Therefore, we should test this op's behavior when used with `.skip()`.
+
+    Args:
+      verify_fn: Verify the correctness of this dataset's checkpointing.
+      symbolic_checkpoint: Whether symbolic checkpointing is turned on.
+      num_skips: `.skip(num_skips)`
+    """
+
+    def build_dataset():
+      def my_map(x):
+        if x == 0:
+          return dataset_ops.Dataset.from_tensor_slices([0, 1, 2, 3])
+        elif x == 1:
+          return dataset_ops.Dataset.from_tensor_slices([4, 5, 6, 7])
+        else:
+          return dataset_ops.Dataset.from_tensor_slices([8, 9, 10, 11])
+
+      indices = dataset_ops.Dataset.from_tensor_slices([0, 1, 2])
+      dataset = indices.flat_map(my_map)
+      # Skip some elements
+      dataset = dataset.skip(num_skips)
+
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(self, build_dataset, num_outputs=3 * 4 - num_skips)
+
+
+@unittest.skip(
+    "TODO: b/355241367 - `flat_map_dataset_op.cc` still needs to be fixed."
+    " Please use concatenate dataset op plus global shuffling instead."
+)
+class FlatMapGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 2],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def test(
+      self,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(1, 10)) * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(dataset_output, self.evaluate(dataset.cardinality()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInputCardinalityTooLarge(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([[i] for i in range(101)])
+    dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        "The cardinality of the input to FlatMapDataset is too large to support"
+        " global shuffling",
+    ):
+      dataset = global_shuffle_op._global_shuffle(dataset, seed=42)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+@unittest.skip(
+    "TODO: b/355241367 - `flat_map_dataset_op.cc` still needs to be fixed."
+    " Please use concatenate dataset op plus global shuffling instead."
+)
+class FlatMapGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 2],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def test(
+      self,
+      verify_fn: Callable[..., None],
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.from_tensor_slices(
+          [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=9 * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/from_generator_test.py

@@ -0,0 +1,529 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tf.data.Dataset.from_generator()."""
+import threading
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import from_generator_op
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+class FromGeneratorTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _testFromGenerator(self, generator, elem_sequence, num_repeats,
+                         requires_initialization):
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int64).repeat(num_repeats).prefetch(5)
+    self.assertDatasetProduces(
+        dataset,
+        elem_sequence * num_repeats,
+        requires_initialization=requires_initialization,
+        num_test_iterations=2)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_repeats=[1, 2], requires_initialization=[True, False])))
+  def testFromGeneratorUsingFn(self, num_repeats, requires_initialization):
+
+    def generator():
+      for i in range(1, 100):
+        yield [i] * i
+
+    elem_sequence = list(generator())
+    self._testFromGenerator(
+        generator,
+        elem_sequence,
+        num_repeats=num_repeats,
+        requires_initialization=requires_initialization)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_repeats=[1, 2], requires_initialization=[True, False])))
+  def testFromGeneratorUsingList(self, num_repeats, requires_initialization):
+    generator = lambda: [[i] * i for i in range(1, 100)]
+    elem_sequence = list(generator())
+    self._testFromGenerator(
+        generator,
+        elem_sequence,
+        num_repeats=num_repeats,
+        requires_initialization=requires_initialization)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_repeats=[1, 2], requires_initialization=[True, False])))
+  def testFromGeneratorUsingNdarray(self, num_repeats, requires_initialization):
+    generator = lambda: np.arange(100, dtype=np.int64)
+    elem_sequence = list(generator())
+    self._testFromGenerator(
+        generator,
+        elem_sequence,
+        num_repeats=num_repeats,
+        requires_initialization=requires_initialization)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_repeats=[1, 2], requires_initialization=[True, False])))
+  def testFromGeneratorUsingGeneratorExpression(self, num_repeats,
+                                                requires_initialization):
+    # NOTE(mrry): Generator *expressions* are not repeatable (or in general
+    # reusable), because they eagerly evaluate the `for` expression as
+    # `iter(range(1, 100))` and discard the means of reconstructing
+    # `range(1, 100)`. Wrapping the generator expression in a `lambda` makes
+    # it repeatable.
+    generator = lambda: ([i] * i for i in range(1, 100))
+    elem_sequence = list(generator())
+    self._testFromGenerator(
+        generator,
+        elem_sequence,
+        num_repeats=num_repeats,
+        requires_initialization=requires_initialization)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromMultipleConcurrentGenerators(self):
+    num_inner_repeats = 5
+    num_outer_repeats = 20
+
+    def generator():
+      for i in range(1, 10):
+        yield ([i] * i, [i, i ** 2, i ** 3])
+    input_list = list(generator())
+
+    # The interleave transformation is essentially a flat map that draws from
+    # multiple input datasets concurrently (in a cyclic fashion). By placing
+    # `Dataset.from_generator()` inside an interleave, we test its behavior when
+    # multiple iterators are active at the same time; by additionally
+    # prefetching inside the interleave, we create the possibility of concurrent
+    # invocations to several iterators created by the same dataset.
+    def interleave_fn(_):
+      return (dataset_ops.Dataset.from_generator(
+          generator, output_types=(dtypes.int64, dtypes.int64),
+          output_shapes=([None], [3]))
+              .repeat(num_inner_repeats).prefetch(5))
+
+    dataset = dataset_ops.Dataset.range(num_outer_repeats).interleave(
+        interleave_fn, cycle_length=10, block_length=len(input_list))
+    get_next = self.getNext(dataset)
+    for _ in range(num_inner_repeats * num_outer_repeats):
+      for elem in input_list:
+        val0, val1 = self.evaluate(get_next())
+        self.assertAllEqual(elem[0], val0)
+        self.assertAllEqual(elem[1], val1)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  def DISABLED_testFromGeneratorsRunningInParallel(self):
+    self.skipTest("b/67868766")
+    num_parallel_iterators = 3
+
+    # Define shared state that multiple iterator instances will access to
+    # demonstrate their concurrent activity.
+    lock = threading.Lock()
+    condition = threading.Condition(lock)
+    next_ticket = [0]  # GUARDED_BY(lock)
+
+    def generator():
+      # NOTE(mrry): We yield one element before the barrier, because
+      # the current implementation of `Dataset.interleave()` must
+      # fetch one element from each incoming dataset to start the
+      # prefetching.
+      yield 0
+
+      # Define a barrier that `num_parallel_iterators` iterators must enter
+      # before any can proceed. Demonstrates that multiple iterators may be
+      # active at the same time.
+      condition.acquire()
+      ticket = next_ticket[0]
+      next_ticket[0] += 1
+      if ticket == num_parallel_iterators - 1:
+        # The last iterator to join the barrier notifies the others.
+        condition.notify_all()
+      else:
+        # Wait until the last iterator enters the barrier.
+        while next_ticket[0] < num_parallel_iterators:
+          condition.wait()
+      condition.release()
+
+      yield 1
+
+    # As in `testFromMultipleConcurrentGenerators()`, we use a combination of
+    # `Dataset.interleave()` and `Dataset.prefetch()` to cause multiple
+    # iterators to be active concurrently.
+    def interleave_fn(_):
+      return dataset_ops.Dataset.from_generator(
+          generator, output_types=dtypes.int64, output_shapes=[]).prefetch(2)
+
+    dataset = dataset_ops.Dataset.range(num_parallel_iterators).interleave(
+        interleave_fn, cycle_length=num_parallel_iterators, block_length=1)
+    get_next = self.getNext(dataset)
+
+    for elem in [0, 1]:
+      for _ in range(num_parallel_iterators):
+        self.assertAllEqual(elem, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorImplicitConversion(self):
+    def generator():
+      yield [1]
+      yield [2]
+      yield [3]
+
+    for dtype in [dtypes.int8, dtypes.int32, dtypes.int64]:
+      dataset = dataset_ops.Dataset.from_generator(
+          generator, output_types=dtype, output_shapes=[1])
+      get_next = self.getNext(dataset)
+
+      for expected in [[1], [2], [3]]:
+        next_val = self.evaluate(get_next())
+        self.assertEqual(dtype.as_numpy_dtype, next_val.dtype)
+        self.assertAllEqual(expected, next_val)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorString(self):
+    def generator():
+      yield "foo"
+      yield b"bar"
+      yield u"baz"
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.string, output_shapes=[])
+    self.assertDatasetProduces(
+        dataset, expected_output=[b"foo", b"bar", b"baz"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorDatastructures(self):
+    # Tests multiple datastructures.
+    def generator():
+      yield {"a": "foo", "b": [1, 2], "c": (9,)}
+      yield {"a": "bar", "b": [3], "c": (7, 6)}
+      yield {"a": "baz", "b": [5, 6], "c": (5, 4)}
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator,
+        output_types={"a": dtypes.string, "b": dtypes.int32, "c": dtypes.int32},
+        output_shapes={"a": [], "b": [None], "c": [None]})
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[{"a": b"foo", "b": [1, 2], "c": [9]},
+                         {"a": b"bar", "b": [3], "c": [7, 6]},
+                         {"a": b"baz", "b": [5, 6], "c": [5, 4]}])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorTypeError(self):
+    def generator():
+      yield np.array([1, 2, 3], dtype=np.int64)
+      yield np.array([4, 5, 6], dtype=np.int64)
+      yield "ERROR"
+      yield np.array([7, 8, 9], dtype=np.int64)
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int64, output_shapes=[3])
+
+    get_next = self.getNext(dataset)
+
+    self.assertAllEqual([1, 2, 3], self.evaluate(get_next()))
+    self.assertAllEqual([4, 5, 6], self.evaluate(get_next()))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+    self.assertAllEqual([7, 8, 9], self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorShapeError(self):
+    def generator():
+      yield np.array([1, 2, 3], dtype=np.int64)
+      yield np.array([4, 5, 6], dtype=np.int64)
+      yield np.array([7, 8, 9, 10], dtype=np.int64)
+      yield np.array([11, 12, 13], dtype=np.int64)
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int64, output_shapes=[3])
+    get_next = self.getNext(dataset)
+
+    self.assertAllEqual([1, 2, 3], self.evaluate(get_next()))
+    self.assertAllEqual([4, 5, 6], self.evaluate(get_next()))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+    self.assertAllEqual([11, 12, 13], self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorStructureError(self):
+    def generator():
+      yield 1, 2
+      yield 3, 4
+      yield 5
+      yield 6, 7, 8
+      yield 9, 10
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=(dtypes.int64, dtypes.int64))
+    get_next = self.getNext(dataset)
+
+    self.assertEqual((1, 2), self.evaluate(get_next()))
+    self.assertEqual((3, 4), self.evaluate(get_next()))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+    self.assertEqual((9, 10), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorHeterogeneous(self):
+    def generator():
+      yield 1
+      yield [2, 3]
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int64)
+    self.assertDatasetProduces(dataset, expected_output=[1, [2, 3]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorStopShort(self):
+
+    def generator():
+      yield 0
+      yield 1
+      yield 2
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int64)
+    get_next = self.getNext(dataset)
+    self.assertAllEqual(0, self.evaluate(get_next()))
+    self.assertAllEqual(1, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorDestructorCalled(self):
+    # Use an `Event` to signal that the generator has been deleted.
+    event = threading.Event()
+
+    class GeneratorWrapper:
+
+      def __iter__(self):
+        return self
+
+      def next(self):
+        return self.__next__()
+
+      def __next__(self):
+        return 42
+
+      def __del__(self):
+        event.set()
+
+    dataset = dataset_ops.Dataset.from_generator(
+        GeneratorWrapper, output_types=dtypes.int64).take(2)
+    get_next = self.getNext(dataset)
+
+    self.assertAllEqual(42, self.evaluate(get_next()))
+    self.assertAllEqual(42, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    # Test that `GeneratorWrapper` object is destroyed when the
+    # iterator terminates (and the generator iterator is deleted).
+    self.assertTrue(event.is_set())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorWithArgs(self):
+
+    def flat_map_fn(elem):
+
+      def generator_with_arg(n):
+        for _ in range(n):
+          yield np.array(n, dtype=np.int64)
+
+      return dataset_ops.Dataset.from_generator(
+          generator_with_arg, output_types=dtypes.int64, output_shapes=(),
+          args=(elem,))
+
+    dataset = dataset_ops.Dataset.range(5).flat_map(flat_map_fn)
+    self.assertDatasetProduces(
+        dataset, expected_output=[1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorWithTwoArgs(self):
+
+    def flat_map_fn(elem, message):
+
+      def generator_with_arg(n, msg):
+        for i in range(n):
+          yield i, msg
+
+      return dataset_ops.Dataset.from_generator(
+          generator_with_arg, output_types=(dtypes.int64, dtypes.string),
+          output_shapes=((), ()), args=(elem, message))
+
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.range(5),
+         dataset_ops.Dataset.from_tensors("Hi!").repeat(None)
+        )).flat_map(flat_map_fn)
+
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[(0, b"Hi!"), (0, b"Hi!"), (1, b"Hi!"), (0, b"Hi!"),
+                         (1, b"Hi!"), (2, b"Hi!"), (0, b"Hi!"), (1, b"Hi!"),
+                         (2, b"Hi!"), (3, b"Hi!")])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGeneratorDatasetFinalizeFunctionCalled(self):
+    # NOTE(mrry): This test tests the internal `_GeneratorDataset`,
+    # which affords more control over what the finalize function can do than
+    # the `Dataset.from_generator()` wrapper.
+
+    # Use an `Event` to signal that the generator has been deleted.
+    event = threading.Event()
+
+    def finalize_fn(_):
+      def finalize_py_func():
+        event.set()
+        return 0
+      return script_ops.py_func(finalize_py_func, [], [dtypes.int64],
+                                stateful=True)
+
+    dummy = constant_op.constant(37)
+
+    dataset = from_generator_op._GeneratorDataset(
+        dummy, lambda x: x, lambda x: x, finalize_fn,
+        tensor_spec.TensorSpec((), dtypes.int32))
+
+    dataset = dataset.take(2)
+    get_next = self.getNext(dataset)
+
+    self.assertAllEqual(37, self.evaluate(get_next()))
+    self.assertAllEqual(37, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSharedName(self):
+
+    def generator():
+      for _ in range(10):
+        yield [20]
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=(dtypes.int64))
+    get_next = self.getNext(
+        dataset, requires_initialization=True, shared_name="shared_dataset")
+
+    self.assertAllEqual([20], self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorRaggedTensor(self):
+
+    def generator():
+      yield ragged_factory_ops.constant([[1, 2], [3]])
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator,
+        output_signature=ragged_tensor.RaggedTensorSpec(
+            shape=(2, None), dtype=dtypes.int32))
+    get_next = self.getNext(dataset)
+
+    ret = get_next()
+
+    self.assertIsInstance(ret, ragged_tensor.RaggedTensor)
+    self.assertAllEqual([[1, 2], [3]], ret)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromGeneratorSparseTensor(self):
+
+    def generator():
+      yield sparse_tensor.SparseTensor(
+          indices=[[0, 0], [1, 2]],
+          values=constant_op.constant([1, 2], dtype=dtypes.int64),
+          dense_shape=[3, 4])
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator,
+        output_signature=sparse_tensor.SparseTensorSpec([3, 4], dtypes.int64))
+
+    get_next = self.getNext(dataset)
+
+    ret = get_next()
+
+    self.assertIsInstance(ret, sparse_tensor.SparseTensor)
+    self.assertAllEqual([[1, 0, 0, 0], [0, 0, 2, 0], [0, 0, 0, 0]],
+                        sparse_ops.sparse_tensor_to_dense(ret))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTypeIsListError(self):
+
+    def generator():
+      for _ in range(10):
+        yield [20]
+
+    with self.assertRaisesRegex(
+        TypeError, r"Cannot convert the argument `type_value`: "
+        r"\[tf.int64\] to a TensorFlow DType"):
+      dataset_ops.Dataset.from_generator(
+          generator, output_types=[dtypes.int64])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDimensionIsListError(self):
+
+    def generator():
+      for _ in range(10):
+        yield [20]
+
+    with self.assertRaisesRegex(TypeError,
+                                r"Dimension value must be integer or None"):
+      dataset_ops.Dataset.from_generator(
+          generator, output_types=(dtypes.int64), output_shapes=[[1]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+
+    def generator():
+      yield 42
+
+    dataset_ops.Dataset.from_generator(
+        generator,
+        output_types=(dtypes.int64),
+        output_shapes=[1],
+        name="from_generator")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py

@@ -0,0 +1,191 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.from_sparse_tensor_slices()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class FromSparseTensorSlicesTest(test_base.DatasetTestBase,
+                                 parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(tf_api_version=1, mode=["graph"]),
+          combinations.combine(slices=[[
+              [1., 2., 3.], [1.], [1.], [1., 2.], [], [1., 2.], [], [], []
+          ], [[1., 2.], [], [1., 2.], [1.], [1., 2.], [], [1., 2.]]])))
+  def testFromSparseTensorSlices(self, slices):
+    """Test a dataset based on slices of a `tf.sparse.SparseTensor`."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+    get_next = sparse_tensor.SparseTensor(*iterator.get_next())
+
+    with self.cached_session() as sess:
+      # Test with sparse tensor in the appropriate order.
+      # pylint: disable=g-complex-comprehension
+      indices = np.array(
+          [[i, j] for i in range(len(slices)) for j in range(len(slices[i]))])
+      values = np.array([val for s in slices for val in s])
+      # pylint: enable=g-complex-comprehension
+      dense_shape = np.array([len(slices), max(len(s) for s in slices) + 1])
+      sparse_feed = sparse_tensor.SparseTensorValue(indices, values,
+                                                    dense_shape)
+      sess.run(init_op, feed_dict={st: sparse_feed})
+      for i, s in enumerate(slices):
+        results = sess.run(get_next)
+        self.assertAllEqual(s, results.values)
+        expected_indices = np.array(
+            [[j] for j in range(len(slices[i]))]).reshape([-1, 1])
+        self.assertAllEqual(expected_indices, results.indices)
+        self.assertAllEqual(dense_shape[1:], results.dense_shape)
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(tf_api_version=1, mode=["graph"]),
+          combinations.combine(slices=[[
+              [1., 2., 3.], [1.], [1.], [1., 2.], [], [1., 2.], [], [], []
+          ], [[1., 2.], [], [1., 2.], [1.], [1., 2.], [], [1., 2.]]])))
+  def testFromSparseTensorSlicesInReverse(self, slices):
+    """Test a dataset based on slices of a `tf.sparse.SparseTensor` in reverse order."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+
+    with self.cached_session() as sess:
+      # pylint: disable=g-complex-comprehension
+      indices = np.array(
+          [[i, j] for i in range(len(slices)) for j in range(len(slices[i]))])
+      values = np.array([val for s in slices for val in s])
+      # pylint: enable=g-complex-comprehension
+      dense_shape = np.array([len(slices), max(len(s) for s in slices) + 1])
+      # Test with sparse tensor in the reverse order, which is not
+      # currently supported.
+      reverse_order_indices = indices[::-1, :]
+      reverse_order_values = values[::-1]
+      sparse_feed = sparse_tensor.SparseTensorValue(
+          reverse_order_indices, reverse_order_values, dense_shape)
+      with self.assertRaises(errors.UnimplementedError):
+        sess.run(init_op, feed_dict={st: sparse_feed})
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testEmptySparseTensorSlices(self):
+    """Test a dataset based on slices of an empty `tf.sparse.SparseTensor`."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+    get_next = sparse_tensor.SparseTensor(*iterator.get_next())
+
+    with self.cached_session() as sess:
+      # Test with an empty sparse tensor.
+      empty_indices = np.empty((0, 4), dtype=np.int64)
+      empty_values = np.empty((0,), dtype=np.float64)
+      empty_dense_shape = [0, 4, 37, 9]
+      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices, empty_values,
+                                                    empty_dense_shape)
+      sess.run(init_op, feed_dict={st: sparse_feed})
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testEmptySparseTensorSlicesInvalid(self):
+    """Test a dataset based on invalid `tf.sparse.SparseTensor`."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+
+    with self.cached_session() as sess:
+      # Test with an empty sparse tensor but with non empty values.
+      empty_indices = np.empty((0, 4), dtype=np.int64)
+      non_empty_values = [1, 2, 3, 4]
+      empty_dense_shape = [0, 4, 37, 9]
+      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices,
+                                                    non_empty_values,
+                                                    empty_dense_shape)
+      # Here, we expect the test to fail when running the feed.
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(init_op, feed_dict={st: sparse_feed})
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testEmptySparseTensorSlicesInvalid2(self):
+    """Test a dataset based on invalid `tf.sparse.SparseTensor`."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+
+    with self.cached_session() as sess:
+      # Test with an empty sparse tensor but with non empty values.
+      empty_indices = [[]]
+      empty_values = []
+      dense_shape = [1, 1]
+      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices, empty_values,
+                                                    dense_shape)
+      # Here, we expect the test to fail when running the feed.
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(init_op, feed_dict={st: sparse_feed})
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode=["eager"]))
+  def testFromSparseTensorSlicesError(self):
+    with self.assertRaises(AttributeError):
+      dataset_ops.Dataset.from_sparse_tensor_slices(None)
+
+
+class FromSparseTensorSlicesCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  def _build_sparse_tensor_slice_dataset(self, slices):
+    # pylint: disable=g-complex-comprehension
+    indices = np.array(
+        [[i, j] for i in range(len(slices)) for j in range(len(slices[i]))],
+        dtype=np.int64)
+    values = np.array([val for s in slices for val in s], dtype=np.float64)
+    # pylint: enable=g-complex-comprehension
+    dense_shape = np.array(
+        [len(slices), max(len(s) for s in slices) + 1], dtype=np.int64)
+    sparse_components = sparse_tensor.SparseTensor(indices, values, dense_shape)
+    return dataset_ops.Dataset.from_sparse_tensor_slices(sparse_components)
+
+  @combinations.generate(
+      combinations.times(test_base.v1_only_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    slices = [[1., 2., 3.], [1.], [1.], [1., 2.], [], [1., 2.], [], [], []]
+
+    verify_fn(
+        self,
+        lambda: self._build_sparse_tensor_slice_dataset(slices),
+        num_outputs=9,
+        sparse_tensors=True)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/from_tensor_slices_test.py

@@ -0,0 +1,481 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.from_tensor_slices()."""
+import collections
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+
+
+class FromTensorSlicesTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesEmptyComponent(self):
+    components = ()
+
+    with self.assertRaises(ValueError):
+      dataset_ops.Dataset.from_tensor_slices(components)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlices(self):
+    """Test a dataset that represents the slices from a tuple of tensors."""
+    components = (
+        np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile(
+            np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0])
+    )
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    get_next = self.getNext(dataset)
+
+    self.assertEqual(
+        [c.shape[1:] for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+
+    for i in range(4):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component[i], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      results = self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesDataset(self):
+    dss = [dataset_ops.Dataset.range(10) for _ in range(10)]
+    ds = dataset_ops.Dataset.from_tensor_slices(dss)
+    ds = ds.flat_map(lambda x: x)
+    self.assertDatasetProduces(ds, expected_output=list(range(10)) * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesDatasetOfOrderedDict(self):
+    dss = [dataset_ops.Dataset.range(10).map(
+        lambda x: collections.OrderedDict([("x", x)])) for _ in range(10)]
+    ds = dataset_ops.Dataset.from_tensor_slices(dss)
+    ds = ds.flat_map(lambda x: x)
+    self.assertDatasetProduces(
+        ds,
+        expected_output=[collections.OrderedDict([("x", x)])
+                         for x in list(range(10)) * 10])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesDatasetInFunction(self):
+    dss = [dataset_ops.Dataset.range(10) for _ in range(10)]
+    ds = dataset_ops.Dataset.from_tensors(dss)
+    ds = ds.flat_map(dataset_ops.Dataset.from_tensor_slices)
+    ds = ds.flat_map(lambda x: x)
+    self.assertDatasetProduces(ds, expected_output=list(range(10)) * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesSparse(self):
+    """Test a dataset that represents the slices from a tuple of tensors."""
+    components = (sparse_tensor.SparseTensorValue(
+        indices=np.array([[0, 0], [1, 0], [2, 0]]),
+        values=np.array([0, 0, 0]),
+        dense_shape=np.array([3, 1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1], [2, 2]]),
+                      values=np.array([1, 2, 3]),
+                      dense_shape=np.array([3, 3])))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+
+    self.assertEqual(
+        [tensor_shape.TensorShape(c.dense_shape[1:]) for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+
+    expected = [
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[0]]),
+             values=np.array([1]),
+             dense_shape=np.array([3]))),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[1]]),
+             values=np.array([2]),
+             dense_shape=np.array([3]))),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[2]]),
+             values=np.array([3]),
+             dense_shape=np.array([3]))),
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesMixed(self):
+    """Test a dataset that represents the slices from a tuple of tensors."""
+    components = (np.tile(np.array([[1], [2], [3]]), 20),
+                  np.tile(np.array([[12], [13], [14]]), 22),
+                  np.array([37.0, 38.0, 39.0]),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 0], [2, 0]]),
+                      values=np.array([0, 0, 0]),
+                      dense_shape=np.array([3, 1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1], [2, 2]]),
+                      values=np.array([1, 2, 3]),
+                      dense_shape=np.array([3, 3])))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    get_next = self.getNext(dataset)
+    self.assertEqual([
+        tensor_shape.TensorShape(c.dense_shape[1:])
+        if sparse_tensor.is_sparse(c) else c.shape[1:] for c in components
+    ], [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+
+    expected = [
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[0]]),
+             values=np.array([1]),
+             dense_shape=np.array([3]))),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[1]]),
+             values=np.array([2]),
+             dense_shape=np.array([3]))),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[2]]),
+             values=np.array([3]),
+             dense_shape=np.array([3]))),
+    ]
+    for i in range(3):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(
+          (list(zip(*components[:3]))[i] + expected[i]), results):
+        self.assertValuesEqual(component, result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesWithDict(self):
+    components = {"foo": [1, 2, 3], "bar": [[4.0], [5.0], [6.0]]}
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    get_next = self.getNext(dataset)
+
+    self.assertEqual(dtypes.int32,
+                     dataset_ops.get_legacy_output_types(dataset)["foo"])
+    self.assertEqual(dtypes.float32,
+                     dataset_ops.get_legacy_output_types(dataset)["bar"])
+    self.assertEqual((), dataset_ops.get_legacy_output_shapes(dataset)["foo"])
+    self.assertEqual((1,), dataset_ops.get_legacy_output_shapes(dataset)["bar"])
+
+    for i in range(3):
+      results = self.evaluate(get_next())
+      self.assertEqual(components["foo"][i], results["foo"])
+      self.assertEqual(components["bar"][i], results["bar"])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesRagged(self):
+    components = (
+        ragged_factory_ops.constant_value([[[0]], [[1]], [[2]]]),
+        ragged_factory_ops.constant_value([[[3]], [[4]], [[5]]]),
+    )
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    expected = [(ragged_factory_ops.constant_value([[0]]),
+                 ragged_factory_ops.constant_value([[3]])),
+                (ragged_factory_ops.constant_value([[1]]),
+                 ragged_factory_ops.constant_value([[4]])),
+                (ragged_factory_ops.constant_value([[2]]),
+                 ragged_factory_ops.constant_value([[5]]))]
+    self.assertDatasetProduces(dataset, expected_output=expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesMixedRagged(self):
+    components = (np.tile(np.array([[1], [2], [3]]),
+                          20), np.tile(np.array([[12], [13], [14]]),
+                                       22), np.array([37.0, 38.0, 39.0]),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 0], [2, 0]]),
+                      values=np.array([0, 0, 0]),
+                      dense_shape=np.array([3, 1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1], [2, 2]]),
+                      values=np.array([1, 2, 3]),
+                      dense_shape=np.array([3, 3])),
+                  ragged_factory_ops.constant_value([[[0]], [[1]], [[2]]]))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    get_next = self.getNext(dataset)
+
+    expected = [
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[0]]),
+             values=np.array([1]),
+             dense_shape=np.array([3])), ragged_factory_ops.constant_value([[0]
+                                                                           ])),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[1]]),
+             values=np.array([2]),
+             dense_shape=np.array([3])), ragged_factory_ops.constant_value([[1]
+                                                                           ])),
+        (sparse_tensor.SparseTensorValue(
+            indices=np.array([[0]]),
+            values=np.array([0]),
+            dense_shape=np.array([1])),
+         sparse_tensor.SparseTensorValue(
+             indices=np.array([[2]]),
+             values=np.array([3]),
+             dense_shape=np.array([3])), ragged_factory_ops.constant_value([[2]
+                                                                           ])),
+    ]
+    for i in range(3):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(
+          (list(zip(*components[:3]))[i] + expected[i]), results):
+        self.assertValuesEqual(component, result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesWithUintDtypes(self):
+    components = (
+        np.tile(np.array([[0], [1]], dtype=np.uint8), 2),
+        np.tile(np.array([[2], [256]], dtype=np.uint16), 2),
+        np.tile(np.array([[4], [65536]], dtype=np.uint32), 2),
+        np.tile(np.array([[8], [4294967296]], dtype=np.uint64), 2),
+    )
+    expected_types = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+    expected_output = [tuple([c[i] for c in components]) for i in range(2)]
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    self.assertEqual(expected_types,
+                     dataset_ops.get_legacy_output_types(dataset))
+    self.assertDatasetProduces(dataset, expected_output)
+
+
+class FromTensorSlicesRandomAccessTest(test_base.DatasetTestBase,
+                                       parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidIndex(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, -1))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 3))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOneDimensionalArray(self):
+    tensor = [1, 2, 3]
+    dataset = dataset_ops.Dataset.from_tensor_slices(tensor)
+    for i in range(len(tensor)):
+      results = self.evaluate(random_access.at(dataset, i))
+      self.assertAllEqual(tensor[i], results)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTwoDimensionalArray(self):
+    tensor = [[1, 2], [3, 4]]
+    dataset = dataset_ops.Dataset.from_tensor_slices(tensor)
+    for i in range(2):
+      results = self.evaluate(random_access.at(dataset, i))
+      self.assertAllEqual(tensor[i], results)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleComponents(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(([1, 2], [3, 4], [5, 6]))
+    self.assertEqual((1, 3, 5), self.evaluate(random_access.at(dataset, 0)))
+    self.assertEqual((2, 4, 6), self.evaluate(random_access.at(dataset, 1)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDictionary(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices({"a": [1, 2], "b": [3, 4]})
+    self.assertEqual({
+        "a": 1,
+        "b": 3
+    }, self.evaluate(random_access.at(dataset, 0)))
+    self.assertEqual({
+        "a": 2,
+        "b": 4
+    }, self.evaluate(random_access.at(dataset, 1)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNumpy(self):
+    components = (
+        np.tile(np.array([[0], [1]], dtype=np.uint8), 2),
+        np.tile(np.array([[2], [256]], dtype=np.uint16), 2),
+        np.tile(np.array([[4], [65536]], dtype=np.uint32), 2),
+        np.tile(np.array([[8], [4294967296]], dtype=np.uint64), 2),
+    )
+    expected_output = [tuple([c[i] for c in components]) for i in range(2)]
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    for i in range(2):
+      result = self.evaluate(random_access.at(dataset, i))
+      self.assertAllEqual(expected_output[i], result)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([42],
+                                                     name="from_tensor_slices")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class FromTensorSlicesCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                     parameterized.TestCase):
+
+  def _build_tensor_slices_dataset(self, components, options=None):
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    # Equal length components
+    components = (np.tile(np.array([[1], [2], [3], [4]]),
+                          20), np.tile(np.array([[12], [13], [14], [15]]),
+                                       22), np.array([37.0, 38.0, 39.0, 40.0]))
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_tensor_slices_dataset(components, options),
+        num_outputs=4)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testDict(self, verify_fn):
+    dict_components = {"foo": [1, 2, 3], "bar": [[4.0], [5.0], [6.0]]}
+
+    verify_fn(
+        self,
+        lambda: self._build_tensor_slices_dataset(dict_components),
+        num_outputs=3)
+
+
+class FromTensorSlicesGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10, 100],
+              repetitions=[1, 3],
+              seed=[None, 19],
+              reshuffle_each_iteration=[True, False])))
+  def testGlobalShuffleTensorSlicesDataset(
+      self,
+      dataset_range: int,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.from_tensor_slices(list(range(dataset_range)))
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+
+    expected = list(range(dataset_range)) * repetitions
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(expected, self.evaluate(dataset.cardinality()))
+
+
+class FromTensorSlicesGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              repetitions=[1, 3],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testGlobalShuffleTensorSlicesDataset(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.from_tensor_slices(
+          list(range(dataset_range)))
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      if symbolic_checkpoint:
+        options = options_lib.Options()
+        options.experimental_symbolic_checkpoint = symbolic_checkpoint
+        dataset = dataset.with_options(options)
+      return dataset
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=dataset_range * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/from_tensors_test.py

@@ -0,0 +1,473 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.from_tensors()."""
+import collections
+import dataclasses
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+
+try:
+  import attr  # pylint:disable=g-import-not-at-top
+except ImportError:
+  attr = None
+
+
+@dataclasses.dataclass
+class MaskedTensor:
+  mask: bool
+  value: np.ndarray
+
+  def __tf_flatten__(self):
+    metadata = (self.mask,)
+    components = (self.value,)
+    return metadata, components
+
+  def __tf_unflatten__(self, metadata, components):
+    mask = metadata[0]
+    value = components[0]
+    return MaskedTensor(mask=mask, value=value)
+
+
+class FromTensorsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensors(self):
+    """Test a dataset that represents a single tuple of tensors."""
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+
+    self.assertEqual(
+        [c.shape for c in components],
+        nest.flatten(dataset_ops.get_legacy_output_shapes(dataset)))
+
+    self.assertDatasetProduces(dataset, expected_output=[components])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsDataset(self):
+    """Test a dataset that represents a dataset."""
+    dataset = dataset_ops.Dataset.from_tensors(dataset_ops.Dataset.range(10))
+    dataset = dataset.flat_map(lambda x: x)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsTensorArray(self):
+    """Test a dataset that represents a TensorArray."""
+    components = (
+        tensor_array_ops.TensorArray(dtypes.float32, element_shape=(), size=2)
+        .unstack([1.0, 2.0]))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+
+    self.assertDatasetProduces(
+        dataset, expected_output=[[1.0, 2.0]], requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsSparse(self):
+    """Test a dataset that represents a single tuple of tensors."""
+    components = (sparse_tensor.SparseTensorValue(
+        indices=np.array([[0]]),
+        values=np.array([0]),
+        dense_shape=np.array([1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1]]),
+                      values=np.array([-1, 1]),
+                      dense_shape=np.array([2, 2])))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+
+    self.assertEqual(
+        [tensor_shape.TensorShape(c.dense_shape) for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    self.assertDatasetProduces(dataset, expected_output=[components])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsMixed(self):
+    """Test an dataset that represents a single tuple of tensors."""
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0]]),
+                      values=np.array([0]),
+                      dense_shape=np.array([1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1]]),
+                      values=np.array([-1, 1]),
+                      dense_shape=np.array([2, 2])))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    self.assertEqual([
+        tensor_shape.TensorShape(c.dense_shape)
+        if sparse_tensor.is_sparse(c) else c.shape for c in components
+    ], [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+
+    self.assertDatasetProduces(dataset, expected_output=[components])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsRagged(self):
+    components = (
+        ragged_factory_ops.constant_value([[[0]], [[1]], [[2]]]),
+        ragged_factory_ops.constant_value([[[3]], [[4]], [[5]]]),
+    )
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+
+    self.assertDatasetProduces(dataset, expected_output=[components])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsNamedTuple(self):
+    Foo = collections.namedtuple("Foo", ["x", "y"])
+    element = Foo(x=1, y=2)
+    dataset = dataset_ops.Dataset.from_tensors(element)
+    self.assertDatasetProduces(dataset, expected_output=[element])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsAttrs(self):
+    if attr is None:
+      self.skipTest("attr module is not available.")
+
+    @attr.s
+    class Foo:
+      x = attr.ib()
+      y = attr.ib()
+
+    element = Foo(x=1, y=2)
+    dataset = dataset_ops.Dataset.from_tensors(element)
+    self.assertDatasetProduces(dataset, expected_output=[element])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsDataclass(self):
+    mt = MaskedTensor(mask=True, value=np.array([1]))
+    dataset = dataset_ops.Dataset.from_tensors(mt)
+    self.assertDatasetProduces(dataset, expected_output=[mt])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsMixedRagged(self):
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0]]),
+                      values=np.array([0]),
+                      dense_shape=np.array([1])),
+                  sparse_tensor.SparseTensorValue(
+                      indices=np.array([[0, 0], [1, 1]]),
+                      values=np.array([-1, 1]),
+                      dense_shape=np.array([2, 2])),
+                  ragged_factory_ops.constant_value([[[0]], [[1]], [[2]]]))
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+
+    self.assertDatasetProduces(dataset, expected_output=[components])
+
+  @combinations.generate(
+      combinations.combine(
+          tf_api_version=[1],
+          mode=["graph"],
+          components=(np.array([1, 2, 3], dtype=np.int64),
+                      (np.array([4., 5.]), np.array(
+                          [6., 7.])), np.array([8, 9, 10], dtype=np.int64)),
+          expected_shapes=[[[None, 3], [None, 3], [None, 2], [None, 2]]]) +
+      combinations.combine(
+          tf_api_version=[1],
+          mode=["eager"],
+          components=(np.array([1, 2, 3], dtype=np.int64),
+                      (np.array([4., 5.]), np.array(
+                          [6., 7.])), np.array([8, 9, 10], dtype=np.int64)),
+          expected_shapes=[[[1, 3], [1, 3], [1, 2], [1, 2]]]))
+  def testNestedStructure(self, components, expected_shapes):
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    dataset = dataset.map(lambda x, y, z: ((x, z), (y[0], y[1])))
+
+    dataset = dataset.flat_map(
+        lambda x, y: dataset_ops.Dataset.from_tensors(
+            ((x[0], x[1]), (y[0], y[1])))).batch(32)
+
+    get_next = self.getNext(dataset)
+    (w, x), (y, z) = get_next()
+    self.assertEqual(dtypes.int64, w.dtype)
+    self.assertEqual(dtypes.int64, x.dtype)
+    self.assertEqual(dtypes.float64, y.dtype)
+    self.assertEqual(dtypes.float64, z.dtype)
+    self.assertEqual(expected_shapes, [
+        w.shape.as_list(),
+        x.shape.as_list(),
+        y.shape.as_list(),
+        z.shape.as_list()
+    ])
+
+    get_next = self.getNext(dataset)
+    (w, x), (y, z) = get_next()
+    self.assertEqual(dtypes.int64, w.dtype)
+    self.assertEqual(dtypes.int64, x.dtype)
+    self.assertEqual(dtypes.float64, y.dtype)
+    self.assertEqual(dtypes.float64, z.dtype)
+    self.assertEqual(expected_shapes, [
+        w.shape.as_list(),
+        x.shape.as_list(),
+        y.shape.as_list(),
+        z.shape.as_list()
+    ])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedDict(self):
+    components = {"a": {"aa": 1, "ab": [2.0, 2.0]}, "b": [3, 3, 3]}
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    self.assertEqual(dtypes.int32,
+                     dataset_ops.get_legacy_output_types(dataset)["a"]["aa"])
+    self.assertEqual(dtypes.float32,
+                     dataset_ops.get_legacy_output_types(dataset)["a"]["ab"])
+    self.assertEqual(dtypes.int32,
+                     dataset_ops.get_legacy_output_types(dataset)["b"])
+    self.assertEqual([],
+                     dataset_ops.get_legacy_output_shapes(dataset)["a"]["aa"])
+    self.assertEqual([2],
+                     dataset_ops.get_legacy_output_shapes(dataset)["a"]["ab"])
+    self.assertEqual([3],
+                     dataset_ops.get_legacy_output_shapes(dataset)["b"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNonSequenceNestedStructure(self):
+    components = np.array([1, 2, 3], dtype=np.int64)
+
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    self.assertEqual(dtypes.int64,
+                     dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual([3], dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.filter(
+        lambda x: math_ops.reduce_all(math_ops.equal(x, components)))
+    self.assertEqual(dtypes.int64,
+                     dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual([3], dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.map(lambda x: array_ops_stack.stack([x, x]))
+    self.assertEqual(dtypes.int64,
+                     dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual([2, 3], dataset_ops.get_legacy_output_shapes(dataset))
+
+    dataset = dataset.flat_map(
+        lambda x: dataset_ops.Dataset.from_tensor_slices(x))
+    self.assertEqual(dtypes.int64,
+                     dataset_ops.get_legacy_output_types(dataset))
+    self.assertEqual([3], dataset_ops.get_legacy_output_shapes(dataset))
+
+    get_next = self.getNext(dataset)
+    self.assertEqual(dtypes.int64, get_next().dtype)
+    self.assertEqual([3], get_next().shape)
+
+  # TODO(b/121264236): needs mechanism for multiple device in eager mode.
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSplitPipeline(self):
+    with session.Session(
+        target="",
+        config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
+
+      dataset = dataset_ops.Dataset.from_tensors(0)
+
+      # Define a pipeline that attempts to use variables on two
+      # different devices.
+      #
+      # Initialize the variables before creating to iterator, to avoid the
+      # placement algorithm overriding the DT_RESOURCE colocation constraints.
+      with ops.device("/cpu:0"):
+        var_0 = resource_variable_ops.ResourceVariable(initial_value=1)
+      dataset = dataset.map(lambda x: x + var_0.read_value())
+      sess.run(var_0.initializer)
+
+      with ops.device("/cpu:1"):
+        var_1 = resource_variable_ops.ResourceVariable(initial_value=1)
+      dataset = dataset.map(lambda x: x + var_1.read_value())
+      sess.run(var_1.initializer)
+
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+      sess.run(iterator.initializer)
+
+      self.assertEqual(sess.run(iterator.get_next()), 2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42, name="from_tensors")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class FromTensorsCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  def _build_tensor_dataset(self, variable_array, options=None):
+    components = (variable_array, np.array([1, 2, 3]), np.array(37.0))
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    arr = np.array(1)
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self, lambda: self._build_tensor_dataset(arr, options), num_outputs=1)
+
+
+class FromTensorsRandomAccessTest(test_base.DatasetTestBase,
+                                  parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidIndex(self):
+    dataset = dataset_ops.Dataset.from_tensors([1, 2, 3])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, -1))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testBasic(self):
+    dataset = dataset_ops.Dataset.from_tensors(range(4))
+    self.assertAllEqual(self.evaluate(random_access.at(dataset, 0)), range(4))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testWithOptions(self):
+    dataset = dataset_ops.Dataset.from_tensors(range(4))
+    options = options_lib.Options()
+    options.experimental_optimization.map_and_batch_fusion = True
+    dataset = dataset.with_options(options)
+
+    self.assertAllEqual(self.evaluate(random_access.at(dataset, 0)), range(4))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.from_tensors([])
+    self.assertAllEqual(self.evaluate(random_access.at(dataset, 0)), [])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testNumpyArray(self):
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    result = self.evaluate(random_access.at(dataset, 0))
+    for i in range(3):
+      self.assertAllEqual(result[i], components[i])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsNestedDataset(self):
+    dataset = dataset_ops.Dataset.from_tensors(dataset_ops.Dataset.range(10))
+    result = random_access.at(dataset, 0)
+    for i in range(10):
+      self.assertEqual(self.evaluate(random_access.at(result, i)), i)
+
+
+class FromTensorsGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 3],
+              seed=[None, 19],
+              reshuffle_each_iteration=[True, False])))
+  def testTensorsDataset(
+      self,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+    dataset = dataset_ops.Dataset.from_tensors(components)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = [components] * repetitions
+    self.assertDatasetProduces(
+        dataset, expected_output=expected, requires_initialization=True)
+    self.assertLen(expected, self.evaluate(dataset.cardinality()))
+
+
+class FromTensorsGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 3],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testGlobalShuffleTensorsDataset(
+      self,
+      verify_fn: Callable[..., None],
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.from_tensors(dataset_ops.Dataset.range(10))
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      dataset = dataset.flat_map(lambda x: x)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=10 * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/get_single_element_test.py

@@ -0,0 +1,134 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the experimental input pipeline ops."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class GetSingleElementTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              skip=[0, 5, 10], take=[1], error=[None], error_msg=[None]) +
+          combinations.combine(
+              skip=[100],
+              take=[1],
+              error=[errors.InvalidArgumentError],
+              error_msg=["Dataset was empty."]) + combinations.combine(
+                  skip=[0],
+                  take=[2],
+                  error=[errors.InvalidArgumentError],
+                  error_msg=["Dataset had more than one element."])))
+  def testBasic(self, skip, take, error=None, error_msg=None):
+
+    def make_sparse(x):
+      x_1d = array_ops.reshape(x, [1])
+      x_2d = array_ops.reshape(x, [1, 1])
+      return sparse_tensor.SparseTensor(x_2d, x_1d, x_1d)
+
+    dataset = dataset_ops.Dataset.range(100).skip(skip).map(
+        lambda x: (x * x, make_sparse(x))).take(take)
+    if error is None:
+      dense_val, sparse_val = self.evaluate(dataset.get_single_element())
+      self.assertEqual(skip * skip, dense_val)
+      self.assertAllEqual([[skip]], sparse_val.indices)
+      self.assertAllEqual([skip], sparse_val.values)
+      self.assertAllEqual([skip], sparse_val.dense_shape)
+    else:
+      with self.assertRaisesRegex(error, error_msg):
+        self.evaluate(dataset.get_single_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindow(self):
+    """Test that `get_single_element()` can consume a nested dataset."""
+
+    def flat_map_func(ds):
+      batched = ds.batch(2)
+      element = batched.get_single_element()
+      return dataset_ops.Dataset.from_tensors(element)
+
+    dataset = dataset_ops.Dataset.range(10).window(2).flat_map(flat_map_func)
+    self.assertDatasetProduces(dataset,
+                               [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSideEffect(self):
+    counter_var = variables.Variable(0)
+
+    def increment_fn(x):
+      counter_var.assign_add(1)
+      return x
+
+    def dataset_fn():
+      return dataset_ops.Dataset.range(1).map(increment_fn)
+
+    @def_function.function
+    def fn():
+      _ = dataset_fn().get_single_element()
+      return "hello"
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    self.assertEqual(self.evaluate(counter_var), 1)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAutomaticControlDependencies(self):
+    counter_var = variables.Variable(1)
+
+    def increment_fn(x):
+      counter_var.assign(counter_var + 1)
+      return x
+
+    def multiply_fn(x):
+      counter_var.assign(counter_var * 2)
+      return x
+
+    def dataset1_fn():
+      return dataset_ops.Dataset.range(1).map(increment_fn)
+
+    def dataset2_fn():
+      return dataset_ops.Dataset.range(1).map(multiply_fn)
+
+    @def_function.function
+    def fn():
+      _ = dataset1_fn().get_single_element()
+      _ = dataset2_fn().get_single_element()
+      return "hello"
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    self.assertEqual(self.evaluate(counter_var), 4)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42)
+    self.assertEqual(
+        self.evaluate(dataset.get_single_element(name="get_single_element")),
+        42)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/group_by_window_test.py

@@ -0,0 +1,452 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.experimental.group_by_window()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.lib.core import error_codes_pb2
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+
+
+# NOTE(mrry): These tests are based on the tests in bucket_ops_test.py.
+# Currently, they use a constant batch size, though should be made to use a
+# different batch size per key.
+class GroupByWindowTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _dynamicPad(self, bucket, window, window_size):
+    # TODO(mrry): To match `tf.contrib.training.bucket()`, implement a
+    # generic form of padded_batch that pads every component
+    # dynamically and does not rely on static shape information about
+    # the arguments.
+    return dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.from_tensors(bucket),
+         window.padded_batch(
+             32, (tensor_shape.TensorShape([]), tensor_shape.TensorShape(
+                 [None]), tensor_shape.TensorShape([3])))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSingleBucket(self):
+
+    def _map_fn(v):
+      return (v, array_ops.fill([v],
+                                v), array_ops.fill([3],
+                                                   string_ops.as_string(v)))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(
+        math_ops.range(32)).map(_map_fn)
+
+    bucketed_dataset = input_dataset.group_by_window(
+        key_func=lambda x, y, z: 0,
+        reduce_func=lambda k, bucket: self._dynamicPad(k, bucket, 32),
+        window_size=32)
+    get_next = self.getNext(bucketed_dataset)
+
+    which_bucket, bucketed_values = self.evaluate(get_next())
+
+    self.assertEqual(0, which_bucket)
+
+    expected_scalar_int = np.arange(32, dtype=np.int64)
+    expected_unk_int64 = np.zeros((32, 31)).astype(np.int64)
+    for i in range(32):
+      expected_unk_int64[i, :i] = i
+    expected_vec3_str = np.vstack(3 * [np.arange(32).astype(bytes)]).T
+
+    self.assertAllEqual(expected_scalar_int, bucketed_values[0])
+    self.assertAllEqual(expected_unk_int64, bucketed_values[1])
+    self.assertAllEqual(expected_vec3_str, bucketed_values[2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEvenOddBuckets(self):
+
+    def _map_fn(v):
+      return (v, array_ops.fill([v],
+                                v), array_ops.fill([3],
+                                                   string_ops.as_string(v)))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(
+        math_ops.range(64)).map(_map_fn)
+
+    bucketed_dataset = input_dataset.group_by_window(
+        key_func=lambda x, y, z: math_ops.cast(x % 2, dtypes.int64),
+        reduce_func=lambda k, bucket: self._dynamicPad(k, bucket, 32),
+        window_size=32)
+
+    get_next = self.getNext(bucketed_dataset)
+
+    # Get two minibatches (one containing even values, one containing odds)
+    which_bucket_even, bucketed_values_even = self.evaluate(get_next())
+    which_bucket_odd, bucketed_values_odd = self.evaluate(get_next())
+
+    # Count number of bucket_tensors.
+    self.assertEqual(3, len(bucketed_values_even))
+    self.assertEqual(3, len(bucketed_values_odd))
+
+    # Ensure bucket 0 was used for all minibatch entries.
+    self.assertAllEqual(0, which_bucket_even)
+    self.assertAllEqual(1, which_bucket_odd)
+
+    # Test the first bucket outputted, the events starting at 0
+    expected_scalar_int = np.arange(0, 32 * 2, 2, dtype=np.int64)
+    expected_unk_int64 = np.zeros((32, 31 * 2)).astype(np.int64)
+    for i in range(0, 32):
+      expected_unk_int64[i, :2 * i] = 2 * i
+      expected_vec3_str = np.vstack(3 *
+                                    [np.arange(0, 32 * 2, 2).astype(bytes)]).T
+
+    self.assertAllEqual(expected_scalar_int, bucketed_values_even[0])
+    self.assertAllEqual(expected_unk_int64, bucketed_values_even[1])
+    self.assertAllEqual(expected_vec3_str, bucketed_values_even[2])
+
+    # Test the second bucket outputted, the odds starting at 1
+    expected_scalar_int = np.arange(1, 32 * 2 + 1, 2, dtype=np.int64)
+    expected_unk_int64 = np.zeros((32, 31 * 2 + 1)).astype(np.int64)
+    for i in range(0, 32):
+      expected_unk_int64[i, :2 * i + 1] = 2 * i + 1
+      expected_vec3_str = np.vstack(
+          3 * [np.arange(1, 32 * 2 + 1, 2).astype(bytes)]).T
+
+    self.assertAllEqual(expected_scalar_int, bucketed_values_odd[0])
+    self.assertAllEqual(expected_unk_int64, bucketed_values_odd[1])
+    self.assertAllEqual(expected_vec3_str, bucketed_values_odd[2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEvenOddBucketsFilterOutAllOdd(self):
+
+    def _map_fn(v):
+      return {
+          "x": v,
+          "y": array_ops.fill([v], v),
+          "z": array_ops.fill([3], string_ops.as_string(v))
+      }
+
+    def _dynamic_pad_fn(bucket, window, _):
+      return dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.from_tensors(bucket),
+           window.padded_batch(
+               32, {
+                   "x": tensor_shape.TensorShape([]),
+                   "y": tensor_shape.TensorShape([None]),
+                   "z": tensor_shape.TensorShape([3])
+               })))
+
+    input_dataset = dataset_ops.Dataset.from_tensor_slices(math_ops.range(
+        128)).map(_map_fn).filter(lambda d: math_ops.equal(d["x"] % 2, 0))
+
+    bucketed_dataset = input_dataset.group_by_window(
+        key_func=lambda d: math_ops.cast(d["x"] % 2, dtypes.int64),
+        reduce_func=lambda k, bucket: _dynamic_pad_fn(k, bucket, 32),
+        window_size=32)
+
+    get_next = self.getNext(bucketed_dataset)
+
+    # Get two minibatches ([0, 2, ...] and [64, 66, ...])
+    which_bucket0, bucketed_values_even0 = self.evaluate(get_next())
+    which_bucket1, bucketed_values_even1 = self.evaluate(get_next())
+
+    # Ensure that bucket 1 was completely filtered out
+    self.assertAllEqual(0, which_bucket0)
+    self.assertAllEqual(0, which_bucket1)
+    self.assertAllEqual(
+        np.arange(0, 64, 2, dtype=np.int64), bucketed_values_even0["x"])
+    self.assertAllEqual(
+        np.arange(64, 128, 2, dtype=np.int64), bucketed_values_even1["x"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDynamicWindowSize(self):
+    components = np.arange(100).astype(np.int64)
+
+    # Key fn: even/odd
+    # Reduce fn: batches of 5
+    # Window size fn: even=5, odd=10
+
+    def window_size_func(key):
+      window_sizes = constant_op.constant([5, 10], dtype=dtypes.int64)
+      return window_sizes[key]
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.group_by_window(
+        key_func=lambda x: x % 2,
+        reduce_func=lambda _, xs: xs.batch(20),
+        window_size=None,
+        window_size_func=window_size_func)
+
+    get_next = self.getNext(dataset)
+    with self.assertRaises(errors.OutOfRangeError):
+      batches = 0
+      while True:
+        result = self.evaluate(get_next())
+        is_even = all(x % 2 == 0 for x in result)
+        is_odd = all(x % 2 == 1 for x in result)
+        self.assertTrue(is_even or is_odd)
+        expected_batch_size = 5 if is_even else 10
+        self.assertEqual(expected_batch_size, result.shape[0])
+        batches += 1
+
+    self.assertEqual(batches, 15)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimple(self):
+    components = np.random.randint(100, size=(200,)).astype(np.int64)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+        lambda x: x * x)
+    dataset = dataset.group_by_window(
+        key_func=lambda x: x % 2,
+        reduce_func=lambda _, xs: xs.batch(4),
+        window_size=4)
+    get_next = self.getNext(dataset)
+    counts = []
+    with self.assertRaises(errors.OutOfRangeError):
+      while True:
+        result = self.evaluate(get_next())
+        self.assertTrue(
+            all(x % 2 == 0 for x in result) or all(x % 2 == 1) for x in result)
+        counts.append(result.shape[0])
+
+    self.assertEqual(len(components), sum(counts))
+    num_full_batches = len([c for c in counts if c == 4])
+    self.assertGreaterEqual(num_full_batches, 24)
+    self.assertTrue(all(c == 4 for c in counts[:num_full_batches]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImmediateOutput(self):
+    components = np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 2, 2, 0, 0],
+                          dtype=np.int64)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.repeat(-1)
+    dataset = dataset.group_by_window(
+        key_func=lambda x: x % 3,
+        reduce_func=lambda _, xs: xs.batch(4),
+        window_size=4)
+    get_next = self.getNext(dataset)
+    # The input is infinite, so this test demonstrates that:
+    # 1. We produce output without having to consume the entire input,
+    # 2. Different buckets can produce output at different rates, and
+    # 3. For deterministic input, the output is deterministic.
+    for _ in range(3):
+      self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
+      self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next()))
+      self.assertAllEqual([2, 2, 2, 2], self.evaluate(get_next()))
+      self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSmallGroups(self):
+    components = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0], dtype=np.int64)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.group_by_window(
+        key_func=lambda x: x % 2,
+        reduce_func=lambda _, xs: xs.batch(4),
+        window_size=4)
+    get_next = self.getNext(dataset)
+    self.assertAllEqual([0, 0, 0, 0], self.evaluate(get_next()))
+    self.assertAllEqual([1, 1, 1, 1], self.evaluate(get_next()))
+    # The small outputs at the end are deterministically produced in key
+    # order.
+    self.assertAllEqual([0, 0, 0], self.evaluate(get_next()))
+    self.assertAllEqual([1], self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmpty(self):
+    dataset = dataset_ops.Dataset.range(4).group_by_window(
+        key_func=lambda _: 0, reduce_func=lambda _, xs: xs, window_size=0)
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Window size must be greater than zero, but got 0."):
+      print(self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConsumeWindowDatasetMoreThanOnce(self):
+    components = np.random.randint(50, size=(200,)).astype(np.int64)
+
+    def reduce_func(key, window):
+      # Apply two different kinds of padding to the input: tight
+      # padding, and quantized (to a multiple of 10) padding.
+      return dataset_ops.Dataset.zip((
+          window.padded_batch(
+              4, padded_shapes=tensor_shape.TensorShape([None])),
+          window.padded_batch(
+              4, padded_shapes=ops.convert_to_tensor([(key + 1) * 10])),
+      ))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.map(
+        lambda x: array_ops.fill([math_ops.cast(x, dtypes.int32)], x))
+    # pylint: disable=g-long-lambda
+    dataset = dataset.group_by_window(
+        key_func=lambda x: math_ops.cast(
+            array_ops.shape(x)[0] // 10, dtypes.int64),
+        reduce_func=reduce_func,
+        window_size=4)
+
+    get_next = self.getNext(dataset)
+    counts = []
+    with self.assertRaises(errors.OutOfRangeError):
+      while True:
+        tight_result, multiple_of_10_result = self.evaluate(get_next())
+        self.assertEqual(0, multiple_of_10_result.shape[1] % 10)
+        self.assertAllEqual(tight_result,
+                            multiple_of_10_result[:, :tight_result.shape[1]])
+        counts.append(tight_result.shape[0])
+    self.assertEqual(len(components), sum(counts))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuit(self):
+
+    dataset = dataset_ops.Dataset.range(10).group_by_window(
+        key_func=lambda x: x,
+        reduce_func=lambda _, window: window.batch(1),
+        window_size=1)
+    self.assertDatasetProduces(
+        dataset, expected_output=[[i] for i in range(10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGroupByWindowWithAutotune(self):
+    dataset = dataset_ops.Dataset.range(1000).group_by_window(
+        key_func=lambda x: x // 10,
+        reduce_func=lambda key, window: dataset_ops.Dataset.from_tensors(key),
+        window_size=4)
+    dataset = dataset.map(lambda x: x + 1, num_parallel_calls=-1)
+    get_next = self.getNext(dataset)
+    self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGroupByWindowCardinality(self):
+    dataset = dataset_ops.Dataset.range(1).repeat().group_by_window(
+        key_func=lambda x: x % 2,
+        reduce_func=lambda key, window: dataset_ops.Dataset.from_tensors(key),
+        window_size=4)
+    self.assertEqual(self.evaluate(dataset.cardinality()), dataset_ops.INFINITE)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(np.int64(42)).group_by_window(
+        key_func=lambda x: x,
+        reduce_func=lambda key, window: window.batch(4),
+        window_size=4,
+        name="group_by_window")
+    self.assertDatasetProduces(dataset, [[42]])
+
+
+class GroupByWindowCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                  parameterized.TestCase):
+
+  def _build_dataset(self, components):
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).repeat(-1)
+    dataset = dataset.group_by_window(
+        key_func=lambda x: x % 3,
+        reduce_func=lambda _, xs: xs.batch(4),
+        window_size=4)
+    return dataset
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test(self):
+    components = np.array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 2, 2, 0, 0],
+                          dtype=np.int64)
+    self.verify_unused_iterator(
+        lambda: self._build_dataset(components),
+        num_outputs=12,
+        verify_exhausted=False)
+    self.verify_multiple_breaks(
+        lambda: self._build_dataset(components),
+        num_outputs=12,
+        verify_exhausted=False)
+    self.verify_reset_restored_iterator(
+        lambda: self._build_dataset(components),
+        num_outputs=12,
+        verify_exhausted=False)
+
+
+class GroupByWindowErrorMessageTest(
+    test_base.DatasetTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReduceFuncError(self):
+    components = np.random.randint(100, size=(200,)).astype(np.int64)
+
+    def my_reduce_func(_, window_dataset):
+      # Introduce an incorrect padded shape that cannot (currently) be
+      # detected at graph construction time.
+      return window_dataset.padded_batch(
+          4,
+          padded_shapes=(
+              tensor_shape.TensorShape([]),
+              constant_op.constant([5], dtype=dtypes.int64) * -1,
+          ),
+      )
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.map(lambda x: (x, ops.convert_to_tensor([x * x])))
+    dataset = dataset.group_by_window(
+        key_func=lambda x, _: x % 2, reduce_func=my_reduce_func, window_size=32
+    )
+    get_next = self.getNext(dataset)
+    with self.assertRaises(errors.InternalError) as error:
+      self.evaluate(get_next())
+
+    msg = str(error.exception)
+    self.assertIn(error_codes_pb2.Code.Name(errors.INVALID_ARGUMENT), msg)
+    self.assertIn(
+        my_reduce_func.__name__,
+        msg,
+        "{} should show up in the error message".format(
+            my_reduce_func.__name__
+        ),
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPropagateUserDefinedFunctionErrorMessage(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([0])
+
+    def a_cool_user_defined_reduce_func(unused_key, window_dataset):
+      it = iter(window_dataset)
+      l = [next(it) for _ in range(2)]  # This causes OutOfRange error
+      return dataset_ops.Dataset.from_tensor_slices(l)
+
+    dataset = dataset.group_by_window(
+        key_func=lambda x: 0,
+        window_size=2,
+        reduce_func=a_cool_user_defined_reduce_func,
+    )
+
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(
+        errors.InternalError,
+        ".*{}.*".format(a_cool_user_defined_reduce_func.__name__),
+        msg=(
+            "The name of user-defined-function should show up in the error"
+            " message"
+        ),
+    ):
+      # Loop over the dataset
+      with self.assertRaises(errors.OutOfRangeError):
+        while True:
+          self.evaluate(get_next())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/ignore_errors_test.py

@@ -0,0 +1,186 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.ignore_errors`."""
+import os
+import sys
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+_NUMPY_RANDOM_SEED = 42
+
+
+class IgnoreErrorsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapIgnoreError(self):
+    components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).map(
+            lambda x: array_ops.check_numerics(x, "message")).ignore_errors())
+    get_next = self.getNext(dataset)
+
+    for x in [1., 2., 3., 5.]:
+      self.assertEqual(x, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIgnoreError_withLogWarning(self):
+    components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).map(
+            lambda x: array_ops.check_numerics(x, "message")).ignore_errors(
+                log_warning=True))
+    get_next = self.getNext(dataset)
+    with self.captureWritesToStream(sys.stderr) as logged:
+      for x in [1., 2., 3.]:
+        self.assertEqual(x, self.evaluate(get_next()))
+      self.assertEqual(5., self.evaluate(get_next()))
+    expected = "Tensor had NaN values"
+    self.assertIn((expected), logged.contents())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelMapIgnoreError(self):
+    components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).map(
+            lambda x: array_ops.check_numerics(x, "message"),
+            num_parallel_calls=2).prefetch(2).ignore_errors())
+    get_next = self.getNext(dataset)
+
+    for x in [1., 2., 3., 5.]:
+      self.assertEqual(x, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadFileIgnoreError(self):
+
+    def write_string_to_file(value, filename):
+      with open(filename, "w") as f:
+        f.write(value)
+
+    filenames = [
+        os.path.join(self.get_temp_dir(), "file_%d.txt" % i) for i in range(5)
+    ]
+    for filename in filenames:
+      write_string_to_file(filename, filename)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(filenames).map(
+            io_ops.read_file, num_parallel_calls=2).prefetch(2).ignore_errors())
+    get_next = self.getNext(dataset)
+
+    # All of the files are present.
+    for filename in filenames:
+      self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Delete one of the files.
+    os.remove(filenames[0])
+
+    # Attempting to read filenames[0] will fail, but ignore_errors()
+    # will catch the error.
+    get_next = self.getNext(dataset)
+    for filename in filenames[1:]:
+      self.assertEqual(compat.as_bytes(filename), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDatasetIgnoreError(self):
+    filenames = []
+    for i in range(5):
+      fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i)
+      filenames.append(fn)
+      writer = python_io.TFRecordWriter(fn)
+      for _ in range(10):
+        writer.write(b"record")
+      writer.close()
+      # Append corrupted data
+      with open(fn, "a") as f:
+        f.write("corrupted data")
+
+    dataset = readers.TFRecordDataset(filenames).ignore_errors()
+    get_next = self.getNext(dataset)
+
+    # All of the files are present.
+    for _ in filenames:
+      for _ in range(10):
+        self.assertEqual(b"record", self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipIgnoreError(self):
+    a = dataset_ops.Dataset.from_tensor_slices([1., 2., 0., 4.])
+    b = a.map(lambda x: array_ops.check_numerics(1. / x, "error"))
+
+    dataset = dataset_ops.Dataset.zip((b, a)).ignore_errors()
+    get_next = self.getNext(dataset)
+
+    for x in [1., 2., 4.]:
+      self.assertEqual((1. / x, x), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCardinality(self):
+    ds = dataset_ops.Dataset.range(10).ignore_errors()
+    self.assertEqual(self.evaluate(ds.cardinality()), dataset_ops.UNKNOWN)
+
+
+class IgnoreErrorsCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                 parameterized.TestCase):
+
+  def _build_ds(self):
+    components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.map(lambda x: array_ops.check_numerics(x, "message"))
+    dataset = dataset.ignore_errors()
+    options = options_lib.Options()
+    options.experimental_external_state_policy = (
+        options_lib.ExternalStatePolicy.IGNORE)
+    return dataset.with_options(options)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    verify_fn(self, self._build_ds, num_outputs=4)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/interleave_test.py

@@ -0,0 +1,695 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.interleave()`."""
+import multiprocessing
+import os
+import sys
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.platform import test
+
+
+def _interleave(lists, cycle_length, block_length, num_parallel_calls=None):
+  """Reference implementation of interleave used for testing.
+
+  Args:
+    lists: a list of lists to interleave
+    cycle_length: the length of the interleave cycle
+    block_length: the length of the interleave block
+    num_parallel_calls: the number of parallel calls
+
+  Yields:
+    Elements of `lists` interleaved in the order determined by `cycle_length`
+    and `block_length`.
+  """
+  num_open = 0
+
+  # `all_iterators` acts as a queue of iterators over each element of `lists`.
+  all_iterators = [iter(l) for l in lists]
+
+  # `open_iterators` are the iterators whose elements are currently being
+  # interleaved.
+  open_iterators = []
+  if cycle_length is None:
+    # The logic here needs to match interleave C++ kernels.
+    cpu_count = multiprocessing.cpu_count()
+    if hasattr(os, "sched_getaffinity"):
+      try:
+        cpu_count = len(os.sched_getaffinity(0))
+      except NotImplementedError:
+        pass
+    if num_parallel_calls is None:
+      cycle_length = cpu_count
+    elif num_parallel_calls == dataset_ops.AUTOTUNE:
+      cycle_length = (cpu_count + 2) // 3
+    else:
+      cycle_length = min(num_parallel_calls, cpu_count)
+
+  for i in range(cycle_length):
+    if all_iterators:
+      open_iterators.append(all_iterators.pop(0))
+      num_open += 1
+    else:
+      open_iterators.append(None)
+
+  while num_open or all_iterators:
+    for i in range(cycle_length):
+      if open_iterators[i] is None:
+        if all_iterators:
+          open_iterators[i] = all_iterators.pop(0)
+          num_open += 1
+        else:
+          continue
+      for _ in range(block_length):
+        try:
+          yield next(open_iterators[i])
+        except StopIteration:
+          open_iterators[i] = None
+          num_open -= 1
+          break
+
+
+def _repeat(values, count):
+  """Produces a list of lists suitable for testing interleave.
+
+  Args:
+    values: for each element `x` the result contains `[x] * x`
+    count: determines how many times to repeat `[x] * x` in the result
+
+  Returns:
+    A list of lists of values suitable for testing interleave.
+  """
+  return [[value] * value for value in np.tile(values, count)]
+
+
+class InterleaveTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              input_values=[[4, 5, 6]],
+              cycle_length=1,
+              block_length=1,
+              expected_elements=[[
+                  4, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 4, 4, 4, 4, 5, 5,
+                  5, 5, 5, 6, 6, 6, 6, 6, 6
+              ]]) + combinations.combine(
+                  input_values=[[4, 5, 6]],
+                  cycle_length=2,
+                  block_length=1,
+                  expected_elements=[[
+                      4, 5, 4, 5, 4, 5, 4, 5, 5, 6, 6, 4, 6, 4, 6, 4, 6, 4, 6,
+                      5, 6, 5, 6, 5, 6, 5, 6, 5, 6, 6
+                  ]]) + combinations.combine(
+                      input_values=[[4, 5, 6]],
+                      cycle_length=2,
+                      block_length=3,
+                      expected_elements=[[
+                          4, 4, 4, 5, 5, 5, 4, 5, 5, 6, 6, 6, 4, 4, 4, 6, 6, 6,
+                          4, 5, 5, 5, 6, 6, 6, 5, 5, 6, 6, 6
+                      ]]) + combinations.combine(
+                          input_values=[[4, 5, 6]],
+                          cycle_length=7,
+                          block_length=2,
+                          expected_elements=[[
+                              4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6, 6, 4, 4, 5, 5, 6,
+                              6, 4, 4, 5, 5, 6, 6, 5, 6, 6, 5, 6, 6
+                          ]]) +
+          combinations.combine(
+              input_values=[[4, 0, 6]],
+              cycle_length=2,
+              block_length=1,
+              expected_elements=[[
+                  4, 4, 6, 4, 6, 4, 6, 6, 4, 6, 4, 6, 4, 4, 6, 6, 6, 6, 6, 6
+              ]])))
+  def testPythonImplementation(self, input_values, cycle_length, block_length,
+                               expected_elements):
+    input_lists = _repeat(input_values, 2)
+
+    for expected, produced in zip(
+        expected_elements, _interleave(input_lists, cycle_length,
+                                       block_length)):
+      self.assertEqual(expected, produced)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              input_values=[np.int64([4, 5, 6])],
+              cycle_length=1,
+              block_length=3,
+              num_parallel_calls=[None, 1]) + combinations.combine(
+                  input_values=[np.int64([4, 5, 6])],
+                  cycle_length=2,
+                  block_length=[1, 3],
+                  num_parallel_calls=[None, 1, 2]) + combinations.combine(
+                      input_values=[np.int64([4, 5, 6])],
+                      cycle_length=7,
+                      block_length=2,
+                      num_parallel_calls=[None, 1, 3, 5, 7]) +
+          combinations.combine(
+              input_values=[np.int64([4, 5, 6, 7])],
+              cycle_length=None,
+              block_length=3,
+              num_parallel_calls=[None, 1]) + combinations.combine(
+                  input_values=[np.int64([]), np.int64([0, 0, 0])],
+                  cycle_length=2,
+                  block_length=3,
+                  num_parallel_calls=[None]) + combinations.combine(
+                      input_values=[np.int64([4, 0, 6])],
+                      cycle_length=2,
+                      block_length=3,
+                      num_parallel_calls=[None, 1, 2])))
+  def testInterleaveDataset(self, input_values, cycle_length, block_length,
+                            num_parallel_calls):
+    count = 2
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_values).repeat(
+        count).interleave(
+            lambda x: dataset_ops.Dataset.from_tensors(x).repeat(x),
+            cycle_length, block_length, num_parallel_calls)
+    expected_output = [
+        element for element in _interleave(
+            _repeat(input_values, count), cycle_length, block_length,
+            num_parallel_calls)
+    ]
+    self.assertDatasetProduces(dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              input_values=[np.float32([1., np.nan, 2., np.nan, 3.])],
+              cycle_length=1,
+              block_length=3,
+              num_parallel_calls=[None, 1]) + combinations.combine(
+                  input_values=[np.float32([1., np.nan, 2., np.nan, 3.])],
+                  cycle_length=2,
+                  block_length=[1, 3],
+                  num_parallel_calls=[None, 1, 2]) + combinations.combine(
+                      input_values=[np.float32([1., np.nan, 2., np.nan, 3.])],
+                      cycle_length=7,
+                      block_length=2,
+                      num_parallel_calls=[None, 1, 3, 5, 7])))
+  def testInterleaveDatasetError(self, input_values, cycle_length, block_length,
+                                 num_parallel_calls):
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_values).map(
+        lambda x: array_ops.check_numerics(x, "message")).interleave(
+            dataset_ops.Dataset.from_tensors, cycle_length, block_length,
+            num_parallel_calls)
+    get_next = self.getNext(dataset)
+
+    for value in input_values:
+      if np.isnan(value):
+        with self.assertRaises(errors.InvalidArgumentError):
+          self.evaluate(get_next())
+      else:
+        self.assertEqual(value, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInterleaveSparse(self):
+
+    def _map_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2])
+
+    def _interleave_fn(x):
+      return dataset_ops.Dataset.from_tensor_slices(
+          sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values))
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn).interleave(
+        _interleave_fn, cycle_length=1)
+    get_next = self.getNext(dataset)
+    for i in range(10):
+      for j in range(2):
+        expected = [i, 0] if j % 2 == 0 else [0, -i]
+        self.assertAllEqual(expected, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              input_values=[np.int64([4, 5, 6])],
+              cycle_length=1,
+              block_length=3,
+              num_parallel_calls=1) + combinations.combine(
+                  input_values=[np.int64([4, 5, 6])],
+                  cycle_length=2,
+                  block_length=[1, 3],
+                  num_parallel_calls=[1, 2]) + combinations.combine(
+                      input_values=[np.int64([4, 5, 6])],
+                      cycle_length=7,
+                      block_length=2,
+                      num_parallel_calls=[1, 3, 5, 7]) + combinations.combine(
+                          input_values=[np.int64([4, 5, 6, 7])],
+                          cycle_length=None,
+                          block_length=3,
+                          num_parallel_calls=1) + combinations.combine(
+                              input_values=[np.int64([4, 0, 6])],
+                              cycle_length=2,
+                              block_length=3,
+                              num_parallel_calls=[1, 2])))
+  def testSloppyInterleaveDataset(self, input_values, cycle_length,
+                                  block_length, num_parallel_calls):
+    count = 2
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_values).repeat(
+        count).interleave(
+            lambda x: dataset_ops.Dataset.from_tensors(x).repeat(x),
+            cycle_length, block_length, num_parallel_calls)
+    options = options_lib.Options()
+    options.deterministic = False
+    dataset = dataset.with_options(options)
+    expected_output = [
+        element for element in _interleave(
+            _repeat(input_values, count), cycle_length, block_length,
+            num_parallel_calls)
+    ]
+    get_next = self.getNext(dataset)
+    actual_output = []
+    for _ in range(len(expected_output)):
+      actual_output.append(self.evaluate(get_next()))
+    self.assertAllEqual(expected_output.sort(), actual_output.sort())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInterleaveMap(self):
+    dataset = dataset_ops.Dataset.range(100)
+
+    def interleave_fn(x):
+      dataset = dataset_ops.Dataset.from_tensors(x)
+      return dataset.map(lambda x: x + x)
+
+    dataset = dataset.interleave(interleave_fn, cycle_length=5)
+    dataset = dataset.interleave(interleave_fn, cycle_length=5)
+
+    self.assertDatasetProduces(dataset, [4 * x for x in range(100)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelInterleaveCached(self):
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.cache(os.path.join(self.get_temp_dir(), "cache_dir"))
+
+    def interleave_fn(x):
+      return dataset_ops.Dataset.from_tensors(x)
+
+    dataset = dataset.interleave(
+        interleave_fn, cycle_length=2, num_parallel_calls=2)
+    self.assertDatasetProduces(dataset, list(range(5)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              local_determinism=[None, True, False],
+              global_determinism=[True, False])))
+  def testDeterminismConfiguration(self, local_determinism, global_determinism):
+    expect_determinism = local_determinism or (local_determinism is None and
+                                               global_determinism)
+    elements = list(range(1000))
+
+    def dataset_fn(delay_ms):
+
+      def interleave_fn(x):
+        ds = dataset_ops.Dataset.from_tensors(x)
+        if math_ops.equal(x, 0):
+          ds = ds.apply(testing.sleep(delay_ms * 1000))
+        else:
+          ds = ds.apply(testing.sleep(0))
+        return ds
+
+      dataset = dataset_ops.Dataset.from_tensor_slices(elements)
+      dataset = dataset.interleave(
+          interleave_fn,
+          cycle_length=10,
+          num_parallel_calls=10,
+          deterministic=local_determinism)
+      opts = options_lib.Options()
+      opts.deterministic = global_determinism
+      dataset = dataset.with_options(opts)
+      return dataset
+
+    self.checkDeterminism(dataset_fn, expect_determinism, elements)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 1])))
+  def testName(self, num_parallel_calls):
+
+    def fn(x):
+      return dataset_ops.Dataset.from_tensors(x)
+
+    dataset = dataset_ops.Dataset.from_tensors(42).interleave(
+        fn, num_parallel_calls=num_parallel_calls, name="interleave")
+    self.assertDatasetProduces(dataset, [42])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 1])))
+  def testMapFuncMustReturnDataset(self, num_parallel_calls):
+
+    def map_fn(x):
+      return [x]
+
+    with self.assertRaisesRegex(
+        TypeError, "The `map_func` argument must return a `Dataset` object."):
+      dataset_ops.Dataset.from_tensors(42).interleave(
+          map_fn, num_parallel_calls=num_parallel_calls)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 1])))
+  def testMapFuncFailWithErrorContext(self, num_parallel_calls):
+
+    def fn(x):
+      return dataset_ops.Dataset.from_tensors(x // 0)
+
+    dataset = dataset_ops.Dataset.from_tensors(42).interleave(
+        fn, num_parallel_calls=num_parallel_calls, name="interleave")
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        r".*Error in user-defined function passed to .* transformation with "
+        r"iterator: Iterator::Root::.*"):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testSymbolicCheckpointSize(self):
+    if sys.platform == "darwin":
+      self.skipTest(
+          "MacOS does not support symbolic checkpointing."
+      )  # b/284304023
+
+    dataset = dataset_ops.Dataset.range(10)
+    # Each input element to `.interleave` is > 1MB
+    dataset = dataset.map(
+        # Create a huge input element
+        lambda x: stateless_random_ops.stateless_random_uniform(
+            [1_000_000], seed=(42, 42)
+        )
+    )
+    dataset = dataset.interleave(
+        lambda x: dataset_ops.Dataset.range(200),
+        cycle_length=5,
+        num_parallel_calls=None,
+    )
+
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = True
+    dataset = dataset.with_options(options)
+
+    it = dataset.as_numpy_iterator()
+    for _ in range(5):
+      next(it)
+
+    checkpoint = it.save().numpy()
+    self.assertLess(
+        len(checkpoint),
+        5_000,
+        f"The checkpoint should be small enough. Got {len(checkpoint)} bytes",
+    )
+
+
+class InterleaveCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              symbolic_checkpoint=[False, True],
+              cycle_length=2,
+              block_length=[1, 3],
+              num_parallel_calls=[None, 1, 2])))
+  def test(self, verify_fn, symbolic_checkpoint, cycle_length, block_length,
+           num_parallel_calls):
+
+    num_repeats = 2
+    input_values = np.array([2, 3], dtype=np.int64)
+
+    def _build_dataset():
+      dataset = dataset_ops.Dataset.from_tensor_slices(input_values)
+      dataset = dataset.repeat(num_repeats)
+      dataset = dataset.interleave(
+          lambda x: dataset_ops.Dataset.from_tensors(x).repeat(x), cycle_length,
+          block_length, num_parallel_calls)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      dataset = dataset.with_options(options)
+      return dataset
+
+    num_outputs = np.sum(input_values) * num_repeats
+    verify_fn(self, _build_dataset, num_outputs)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              skip=[0, 1, 2, 3],
+          ),
+      )
+  )
+  def testWithSkip(self, verify_fn, skip):
+    def _build_dataset():
+      dataset = dataset_ops.Dataset.range(4)
+
+      dataset = dataset.interleave(
+          lambda x: dataset_ops.Dataset.from_tensors(x).repeat(3),
+          cycle_length=2,
+          block_length=1,
+          num_parallel_calls=None,
+      )
+      dataset = dataset.skip(skip)
+      return dataset
+
+    num_outputs = 4 * 3 - skip
+    verify_fn(self, _build_dataset, num_outputs)
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testDelayedPurgeCheckpointAtTheSameCycleIdx(self):
+    """Tests delayed checkpoint purging at the same cycle index works correctly.
+
+    This would crash if we were to use`cycle_index_` as part
+    of the prefix:
+           [0]                                 [1]
+            1(prefix: ::Interleave[0]           2(prefix: ::Interleave[1])
+           EOF(delete ::Interleave[0])         EOF(delete ::Interleave[1])
+            3(prefix  ::Interleave[2])          4
+                                   ^
+                                  (should be 2 instead of 0)
+           EOF                                 EOF
+
+    If we checkpoint at the point right after 3 is generated and
+    restore it, restore would crash because the sub iterator
+    for generating 3 is incorrectly deleted due to delayed checkpoint purging.
+    """
+
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = True
+    options.experimental_optimization.inject_prefetch = False
+    options.experimental_optimization.apply_default_optimizations = False
+
+    def _build_dataset():
+      dataset = dataset_ops.Dataset.range(4)
+
+      dataset = dataset.interleave(
+          lambda x: dataset_ops.Dataset.from_tensor_slices([x]),
+          cycle_length=2,
+          block_length=1,
+          num_parallel_calls=None,
+      )
+      dataset = dataset.with_options(options)
+      return dataset
+
+    dataset = _build_dataset().with_options(options)
+
+    it = dataset.as_numpy_iterator()
+
+    for _ in range(3):
+      next(it)
+
+    checkpoint = it.save().numpy()
+
+    expected = next(it)
+
+    restored_it = dataset.as_numpy_iterator()
+    restored_it.restore(checkpoint)
+
+    actual = next(restored_it)
+
+    self.assertEqual(expected, actual)
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testWithInputThatPurgeCheckpoint(self):
+    """Tests underlying `expired_prefixes` are handled correctly.
+
+    Explanation:
+        The input for `interleave` looks like (created by `.repeat`):
+        [0, |1, |2]
+            ^   ^
+            |   |
+            |   expired_prefixes=["FiniteRepeat[1]"]
+            expired_prefixes=["FiniteRepeat[0]"]
+
+        [0]   [1]
+         0     1    <--- expired_prefixes=["...FiniteRepeat[0]"]
+        EOF   EOF
+         2 <----- Tests the previous checkpoint stored at this index
+                  should not have an effect on the new checkpoint.
+
+        EOF
+    """
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = True
+    options.experimental_optimization.inject_prefetch = False
+    options.experimental_optimization.apply_default_optimizations = False
+
+    def carefully_designed_map(x):
+      if x == 0:
+        return dataset_ops.Dataset.from_tensor_slices([0])
+      elif x == 1:
+        return dataset_ops.Dataset.from_tensor_slices([1])
+      else:
+        return dataset_ops.Dataset.from_tensor_slices([2])
+
+    def _build_dataset():
+      dataset = dataset_ops.Dataset.from_tensor_slices(["does not matter"])
+
+      # Create [0, 1, 2] using repeat+enumerate+map
+      dataset = dataset.repeat(3)
+      dataset = dataset.enumerate()
+      dataset = dataset.map(lambda idx, x: idx)
+
+      dataset = dataset.interleave(
+          carefully_designed_map,
+          cycle_length=2,
+          block_length=1,
+          num_parallel_calls=None,
+      )
+      dataset = dataset.with_options(options)
+      return dataset
+
+    dataset = _build_dataset().with_options(options)
+
+    it = dataset.as_numpy_iterator()
+
+    try:
+      for _ in range(4):
+        next(it)
+    except StopIteration:
+      pass
+
+    # should not crash
+    it.save().numpy()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(num_parallel_calls=[None, 2]),
+      )
+  )
+  def testNested(self, verify_fn, num_parallel_calls):
+
+    def build_ds():
+
+      inner_ds = dataset_ops.Dataset.from_tensor_slices(range(10))
+      ds = dataset_ops.Dataset.from_tensors(inner_ds).repeat(10)
+      return ds.interleave(
+          lambda x: x, cycle_length=5, num_parallel_calls=num_parallel_calls)
+
+    verify_fn(self, build_ds, num_outputs=100)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testSparse(self, verify_fn):
+
+    def _map_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0], [1, 1]], values=(i * [1, -1]), dense_shape=[2, 2])
+
+    def _interleave_fn(x):
+      return dataset_ops.Dataset.from_tensor_slices(
+          sparse_ops.sparse_to_dense(x.indices, x.dense_shape, x.values))
+
+    def _build_dataset():
+      return dataset_ops.Dataset.range(10).map(_map_fn).interleave(
+          _interleave_fn, cycle_length=1)
+
+    verify_fn(self, _build_dataset, num_outputs=20)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(num_parallel_calls=[None, 2]),
+      )
+  )
+  def testSymbolicUnimplemented(self, verify_fn, num_parallel_calls):
+    if sys.platform == "darwin":
+      self.skipTest(
+          "MacOS does not support symbolic checkpointing."
+      )  # b/284304023
+
+    def fn(x):
+      del x
+      dataset = dataset_ops.Dataset.range(7)
+      dataset = dataset.window(3, shift=1, drop_remainder=True)
+      dataset = dataset.flat_map(lambda x: x)
+      return dataset
+
+    def build_ds():
+      dataset = dataset_ops.Dataset.range(2)
+      dataset = dataset.interleave(
+          fn,
+          cycle_length=3,
+          num_parallel_calls=num_parallel_calls,
+      )
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = True
+      dataset = dataset.with_options(options)
+      return dataset
+
+    with self.assertRaisesRegex(
+        errors.UnimplementedError,
+        "WindowOp does not support symbolic checkpointing.",
+    ):
+      verify_fn(self, build_ds, num_outputs=30)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/io_test.py

@@ -0,0 +1,223 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the `tf.data.experimental.{save,load}` operations."""
+
+import os
+import shutil
+import threading
+import time
+
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class IOTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(IOTest, self).setUp()
+    tmpdir = self.get_temp_dir()
+    tmpdir = os.path.join(tmpdir, "io_test")
+    os.mkdir(tmpdir)
+    self._test_dir = tmpdir
+    self._checkpoint_prefix = os.path.join(self.get_temp_dir(), "ckpt")
+    os.mkdir(self._checkpoint_prefix)
+    self._save_dir = os.path.join(self.get_temp_dir(), "save")
+    os.mkdir(self._save_dir)
+
+  def tearDown(self):
+    super(IOTest, self).tearDown()
+    shutil.rmtree(self._test_dir)
+    shutil.rmtree(self._checkpoint_prefix)
+    shutil.rmtree(self._save_dir)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(compression=[None, "GZIP"]),
+      )
+  )
+  def testBasic(self, compression):
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._test_dir, compression=compression))
+    dataset2 = dataset_ops.Dataset.load(
+        self._test_dir, dataset.element_spec, compression=compression
+    )
+    self.assertDatasetProduces(dataset2, range(42))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCardinality(self):
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._test_dir))
+    dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec)
+    self.assertEqual(self.evaluate(dataset2.cardinality()), 42)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCustomShardFunction(self):
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._test_dir, shard_func=lambda x: x // 21))
+    dataset2 = dataset_ops.Dataset.load(self._test_dir, dataset.element_spec)
+    expected = []
+    for i in range(21):
+      expected.extend([i, i + 21])
+    self.assertDatasetProduces(dataset2, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCustomReaderFunction(self):
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._test_dir, shard_func=lambda x: x % 7))
+    dataset2 = dataset_ops.Dataset.load(
+        self._test_dir,
+        dataset.element_spec,
+        reader_func=lambda x: x.flat_map(lambda y: y),
+    )
+    expected = []
+    for i in range(7):
+      expected.extend(range(i, 42, 7))
+    self.assertDatasetProduces(dataset2, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(compression=[None, "GZIP"]),
+      )
+  )
+  def testSaveInsideFunction(self, compression):
+    dataset = dataset_ops.Dataset.range(42)
+
+    @def_function.function
+    def save_fn():
+      dataset.save(self._test_dir, compression=compression)
+
+    save_fn()
+    dataset = dataset_ops.Dataset.load(
+        self._test_dir, dataset.element_spec, compression=compression
+    )
+    self.assertDatasetProduces(dataset, range(42))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testElementSpecOptional(self):
+    range_dataset = dataset_ops.Dataset.range(42)
+    dict_dataset = dataset_ops.Dataset.from_tensor_slices(
+        {"a": [1, 2], "b": [3, 4]}
+    )
+    tuple_dataset = dataset_ops.Dataset.from_tensor_slices(([1, 2], [3, 4]))
+    dataset = dataset_ops.Dataset.zip(
+        (range_dataset, dict_dataset, tuple_dataset)
+    )
+    self.evaluate(dataset.save(self._test_dir))
+    dataset_loaded = dataset_ops.Dataset.load(self._test_dir)
+    self.assertDatasetsEqual(dataset, dataset_loaded)
+
+  @combinations.generate(test_base.default_test_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))
+    self.evaluate(dataset1.save(self._test_dir))
+    dataset = dataset_ops.Dataset.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())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWait(self):
+    dataset_range = 50
+
+    def load_thread_fn():
+      dataset = dataset_ops.Dataset.load(self._test_dir)
+      self.assertDatasetProduces(
+          dataset, list(range(dataset_range)), assert_items_equal=True)
+    load_thread = threading.Thread(target=load_thread_fn, name="load_thread")
+    load_thread.start()
+
+    def save_thread_fn():
+      time.sleep(5)
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      self.evaluate(dataset.save(self._test_dir))
+    save_thread = threading.Thread(target=save_thread_fn, name="save_thread")
+    save_thread.start()
+    save_thread.join()
+    load_thread.join()
+
+
+class LoadCheckpointTest(IOTest, checkpoint_test_base.CheckpointTestBase):
+
+  def _build_ds(self):
+    return dataset_ops.Dataset.load(self._save_dir)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def test(self, verify_fn):
+    dataset = dataset_ops.Dataset.range(42)
+    self.evaluate(dataset.save(self._save_dir))
+    verify_fn(self, self._build_ds, num_outputs=42)
+
+
+class SaveCheckpointTest(IOTest, checkpoint_test_base.CheckpointTestBase):
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveCheckpointingAPI(self):
+    dataset = dataset_ops.Dataset.range(40)
+    checkpoint_args = {"directory": self._checkpoint_prefix, "max_to_keep": 50}
+    dataset.save(self._save_dir, checkpoint_args=checkpoint_args)
+    num_checkpoint_files = len(list(os.listdir(self._checkpoint_prefix)))
+    # By default, we checkpoint every increment. Each checkpoint writes a
+    # file containing the data and a file containing the index. There is
+    # also an overall checkpoint file. Thus, we expect (2 * 40) + 1 files.
+    self.assertEqual(81, num_checkpoint_files)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveCheckpointingAPICustomCheckpointInterval(self):
+    dataset = dataset_ops.Dataset.range(40)
+    step_counter = variables.Variable(0, trainable=False)
+    checkpoint_args = {
+        "checkpoint_interval": 5,
+        "step_counter": step_counter,
+        "directory": self._checkpoint_prefix,
+        "max_to_keep": 10,
+    }
+    dataset.save(self._save_dir, checkpoint_args=checkpoint_args)
+    num_checkpoint_files = len(list(os.listdir(self._checkpoint_prefix)))
+    # We expect (2 * 8) + 1 files.
+    self.assertEqual(17, num_checkpoint_files)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testSaveCheckpointingAPIIncorrectArgs(self):
+    dataset = dataset_ops.Dataset.range(42)
+    checkpoint_args = {
+        "directory": self._checkpoint_prefix,
+        "incorrect_arg": "incorrect_arg"
+    }
+    with self.assertRaises(TypeError):
+      dataset.save(
+          dataset, self._save_dir, checkpoint_args=checkpoint_args)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/iterator_cluster_test.py

@@ -0,0 +1,180 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Iterator` using distributed sessions."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+
+
+class IteratorClusterTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testRemoteIteratorWithoutRemoteCallFail(self):
+    worker_config = config_pb2.ConfigProto()
+    worker_config.device_count["CPU"] = 2
+    worker, _ = test_util.create_local_cluster(
+        1, 1, worker_config=worker_config)
+
+    with ops.device("/job:worker/replica:0/task:0/cpu:1"):
+      dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
+      iterator_3 = dataset_ops.make_one_shot_iterator(dataset_3)
+      iterator_3_handle = iterator_3.string_handle()
+
+    with ops.device("/job:worker/replica:0/task:0/cpu:0"):
+      remote_it = iterator_ops.Iterator.from_string_handle(
+          iterator_3_handle, dataset_ops.get_legacy_output_types(dataset_3),
+          dataset_ops.get_legacy_output_shapes(dataset_3))
+      get_next_op = remote_it.get_next()
+
+    with session.Session(worker[0].target) as sess:
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(get_next_op)
+
+  def _testRemoteIteratorHelper(self, device0, device1, target):
+    with ops.device(device1):
+      dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
+      iterator_3 = dataset_ops.make_one_shot_iterator(dataset_3)
+      iterator_3_handle = iterator_3.string_handle()
+
+    @function.Defun(dtypes.string)
+    def _remote_fn(h):
+      remote_iterator = iterator_ops.Iterator.from_string_handle(
+          h, dataset_ops.get_legacy_output_types(dataset_3),
+          dataset_ops.get_legacy_output_shapes(dataset_3))
+      return remote_iterator.get_next()
+
+    with ops.device(device0):
+      target_placeholder = array_ops.placeholder(dtypes.string, shape=[])
+      remote_op = functional_ops.remote_call(
+          args=[iterator_3_handle],
+          Tout=[dtypes.int32],
+          f=_remote_fn,
+          target=target_placeholder)
+
+    with session.Session(target) as sess:
+      elem = sess.run(remote_op, feed_dict={target_placeholder: device1})
+      self.assertEqual(elem, [1])
+      # Fails when target is cpu:0 where the resource is not located.
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(remote_op, feed_dict={target_placeholder: device0})
+      elem = sess.run(iterator_3.get_next())
+      self.assertEqual(elem, [2])
+      elem = sess.run(remote_op, feed_dict={target_placeholder: device1})
+      self.assertEqual(elem, [3])
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(remote_op, feed_dict={target_placeholder: device1})
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testRemoteIteratorUsingRemoteCallOp(self):
+    worker_config = config_pb2.ConfigProto()
+    worker_config.device_count["CPU"] = 2
+    worker, _ = test_util.create_local_cluster(
+        1, 1, worker_config=worker_config)
+
+    self._testRemoteIteratorHelper("/job:worker/replica:0/task:0/cpu:0",
+                                   "/job:worker/replica:0/task:0/cpu:1",
+                                   worker[0].target)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testRemoteIteratorUsingRemoteCallOpCrossProcess(self):
+    workers, _ = test_util.create_local_cluster(2, 1)
+
+    self._testRemoteIteratorHelper("/job:worker/replica:0/task:0/cpu:0",
+                                   "/job:worker/replica:0/task:1/cpu:0",
+                                   workers[0].target)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCaptureHashTableInSharedIterator(self):
+    worker, _ = test_util.create_local_cluster(1, 1)
+
+    # NOTE(mrry): We must use the V2 variants of `HashTable`
+    # etc. because these produce a `tf.resource`-typed output that is
+    # compatible with the in-graph function implementation.
+    default_val = -1
+    keys = constant_op.constant(["brain", "salad", "surgery"])
+    values = constant_op.constant([0, 1, 2], dtypes.int64)
+    table = lookup_ops.StaticHashTableV1(
+        lookup_ops.KeyValueTensorInitializer(keys, values),
+        default_val)
+
+    input_sentences = dataset_ops.Dataset.from_tensor_slices(
+        ["brain brain tank salad surgery", "surgery brain"])
+
+    dataset = input_sentences.map(
+        lambda x: string_ops.string_split([x]).values).map(table.lookup)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset, shared_name="shared_iterator")
+    init_op = iterator.initializer
+    get_next = iterator.get_next()
+
+    with session.Session(worker[0].target) as sess:
+      sess.run(table.initializer)
+      sess.run(init_op)
+      self.assertAllEqual([0, 0, -1, 1, 2], sess.run(get_next))
+
+    with session.Session(worker[0].target) as sess:
+      self.assertAllEqual([2, 0], sess.run(get_next))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testImplicitDisposeParallelMapDataset(self):
+    # Tests whether a parallel map dataset will be cleaned up correctly when
+    # the pipeline does not run it until exhaustion.
+    # The pipeline is TensorSliceDataset -> MapDataset(square_3) ->
+    # RepeatDataset(None) -> PrefetchDataset(100).
+    worker, _ = test_util.create_local_cluster(1, 1)
+
+    components = (np.arange(1000),
+                  np.array([[1, 2, 3]]) * np.arange(1000)[:, np.newaxis],
+                  np.array(37.0) * np.arange(1000))
+
+    def _map_fn(x, y, z):
+      return math_ops.square(x), math_ops.square(y), math_ops.square(z)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn)
+        .repeat(None).prefetch(10000))
+
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    init_op = iterator.initializer
+    get_next = iterator.get_next()
+
+    with session.Session(worker[0].target) as sess:
+      sess.run(init_op)
+      for _ in range(3):
+        sess.run(get_next)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/len_test.py

@@ -0,0 +1,58 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.__len__()`."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+class LenTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testKnown(self):
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements)
+    self.assertLen(ds, 10)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testInfinite(self):
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements).repeat()
+    with self.assertRaisesRegex(TypeError, "infinite"):
+      len(ds)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testUnknown(self):
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements).filter(lambda x: True)
+    with self.assertRaisesRegex(TypeError, "unknown"):
+      len(ds)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testGraphMode(self):
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements)
+    with self.assertRaisesRegex(
+        TypeError,
+        r"`tf.data.Dataset` only supports `len` in eager mode. Use "
+        r"`tf.data.Dataset.cardinality\(\)` instead."):
+      len(ds)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/list_files_test.py

@@ -0,0 +1,332 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.list_files()`."""
+
+from os import path
+import shutil
+import tempfile
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import test_mode
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class ListFilesTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(ListFilesTest, self).setUp()
+    self.tmp_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    shutil.rmtree(self.tmp_dir, ignore_errors=True)
+    super(ListFilesTest, self).tearDown()
+
+  def _touchTempFiles(self, filenames):
+    for filename in filenames:
+      open(path.join(self.tmp_dir, filename), 'a').close()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDirectory(self):
+    with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
+                                             'No files matched'):
+      dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'))
+      # We need requires_initialization=True so that getNext uses
+      # make_initializable_iterator instead of make_one_shot_iterator.
+      # make_one_shot_iterator has an issue where it fails to capture control
+      # dependencies when capturing the dataset, so it loses the assertion that
+      # list_files matches at least one file.
+      # TODO(b/140837601): Make this work with make_one_shot_iterator.
+      self.getNext(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimpleDirectory(self):
+    filenames = ['a', 'b', 'c']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimpleDirectoryNotShuffled(self):
+    filenames = ['b', 'c', 'a']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(
+        path.join(self.tmp_dir, '*'), shuffle=False)
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in sorted(filenames)
+        ])
+
+  def testFixedSeedResultsInRepeatableOrder(self):
+    filenames = ['a', 'b', 'c']
+    self._touchTempFiles(filenames)
+
+    def dataset_fn():
+      return dataset_ops.Dataset.list_files(
+          path.join(self.tmp_dir, '*'), shuffle=True, seed=37)
+
+    expected_filenames = [
+        compat.as_bytes(path.join(self.tmp_dir, filename))
+        for filename in filenames
+    ]
+
+    all_actual_filenames = []
+    for _ in range(3):
+      actual_filenames = []
+      next_element = self.getNext(dataset_fn(), requires_initialization=True)
+      try:
+        while True:
+          actual_filenames.append(self.evaluate(next_element()))
+      except errors.OutOfRangeError:
+        pass
+      all_actual_filenames.append(actual_filenames)
+
+    # Each run should produce the same set of filenames, which may be
+    # different from the order of `expected_filenames`.
+    self.assertCountEqual(expected_filenames, all_actual_filenames[0])
+    # However, the different runs should produce filenames in the same order
+    # as each other.
+    self.assertEqual(all_actual_filenames[0], all_actual_filenames[1])
+    self.assertEqual(all_actual_filenames[0], all_actual_filenames[2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def tesEmptyDirectoryInitializer(self):
+
+    def dataset_fn():
+      return dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'))
+
+    self.assertDatasetProduces(
+        dataset_fn(),
+        expected_error=(errors.InvalidArgumentError,
+                        'No files matched pattern'),
+        requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimpleDirectoryInitializer(self):
+    filenames = ['a', 'b', 'c']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileSuffixes(self):
+    filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*.py'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames[1:-1]
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileMiddles(self):
+    filenames = ['a.txt', 'b.py', 'c.pyc']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*.py*'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames[1:]
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoShuffle(self):
+    filenames = ['a', 'b', 'c']
+    self._touchTempFiles(filenames)
+
+    # Repeat the list twice and ensure that the order is the same each time.
+    # NOTE(mrry): This depends on an implementation detail of `list_files()`,
+    # which is that the list of files is captured when the iterator is
+    # initialized. Otherwise, or if e.g. the iterator were initialized more than
+    # once, it's possible that the non-determinism of `tf.matching_files()`
+    # would cause this test to fail. However, it serves as a useful confirmation
+    # that the `shuffle=False` argument is working as intended.
+    # TODO(b/73959787): Provide some ordering guarantees so that this test is
+    # more meaningful.
+    dataset = dataset_ops.Dataset.list_files(
+        path.join(self.tmp_dir, '*'), shuffle=False).repeat(2)
+    next_element = self.getNext(dataset)
+
+    expected_filenames = []
+    actual_filenames = []
+    for filename in filenames * 2:
+      expected_filenames.append(
+          compat.as_bytes(path.join(self.tmp_dir, filename)))
+      actual_filenames.append(compat.as_bytes(self.evaluate(next_element())))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+    self.assertCountEqual(expected_filenames, actual_filenames)
+    self.assertEqual(actual_filenames[:len(filenames)],
+                     actual_filenames[len(filenames):])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultiplePatternsAsList(self):
+    filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc']
+    self._touchTempFiles(filenames)
+
+    patterns = [path.join(self.tmp_dir, pat) for pat in ['*.py', '*.txt']]
+    dataset = dataset_ops.Dataset.list_files(patterns)
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames[:-1]
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultiplePatternsAsTensor(self):
+    filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc']
+    self._touchTempFiles(filenames)
+
+    dataset = dataset_ops.Dataset.list_files(
+        [path.join(self.tmp_dir, pat) for pat in ['*.py', '*.txt']])
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(path.join(self.tmp_dir, filename))
+            for filename in filenames[:-1]
+        ],
+        assert_items_equal=True)
+
+
+class ListFilesGlobalShuffleTest(ListFilesTest, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 2],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def test(
+      self,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    filenames = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k']
+    self._touchTempFiles(filenames)
+    dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'),
+                                             shuffle=False)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = [
+        compat.as_bytes(path.join(self.tmp_dir, filename))
+        for filename in filenames
+    ] * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(dataset_output, self.evaluate(dataset.cardinality()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleNotSupported(self):
+    filenames = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k']
+    self._touchTempFiles(filenames)
+    dataset = dataset_ops.Dataset.list_files(
+        path.join(self.tmp_dir, '*'), shuffle=True)
+    with self.assertRaises(errors.FailedPreconditionError):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class ListFilesGlobalShuffleCheckpointTest(
+    ListFilesTest,
+    checkpoint_test_base.CheckpointTestBase,
+    parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Bypasses the default value for `warm_start`, which is not supported for
+    # global shuffling:
+    # https://github.com/tensorflow/tensorflow/blob/29561af231863afb3b6b89e3aa8a6a550c2b7bb0/tensorflow/python/data/ops/options.py#L633
+    test_mode.toggle_test_mode(False)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 2],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def test(
+      self,
+      verify_fn: Callable[..., None],
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+    filenames = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k']
+    self._touchTempFiles(filenames)
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.list_files(path.join(self.tmp_dir, '*'),
+                                               shuffle=False)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=len(filenames) * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/data/kernel_tests/memory_cleanup_test.py

@@ -0,0 +1,213 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Verify that memory usage is minimal in eager mode."""
+
+import gc
+import time
+import weakref
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import multi_device_iterator_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.types import internal
+
+
+# memory_profiler might not be available in the OSS version of TensorFlow.
+try:
+  import memory_profiler  # pylint:disable=g-import-not-at-top
+except ImportError:
+  memory_profiler = None
+
+
+class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(MemoryCleanupTest, self).setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(3)
+
+  def assertMemoryNotIncreasing(self, f, num_iters, max_increase_mb):
+    """Assert memory usage doesn't increase beyond given threshold for f."""
+
+    # Warm up.
+    f()
+    # Wait for background threads to start up and allocate memory.
+    time.sleep(4)
+    initial = memory_profiler.memory_usage(-1)[0]
+    for _ in range(num_iters):
+      f()
+    increase = memory_profiler.memory_usage(-1)[0] - initial
+    logging.info("Memory increase observed: %f MB" % increase)
+    assert increase < max_increase_mb, (
+        "Increase is too high. Initial memory usage: %f MB. Increase: %f MB. "
+        "Maximum allowed increase: %f") % (initial, increase, max_increase_mb)
+
+  def assertNoMemoryLeak(self, dataset_fn):
+    """Assert consuming elements from the dataset does not leak memory."""
+
+    def run():
+      get_next = self.getNext(dataset_fn())
+      for _ in range(100):
+        self.evaluate(get_next())
+
+    for _ in range(10):
+      run()
+
+    gc.collect()
+
+    def is_native_object(o):
+      # First check if `o` is a weakref proxy. Calling
+      # `isinstance(o, internal.NativeObject)` on an expired weak reference
+      # proxy will raise a ReferenceError.
+      if isinstance(o, weakref.ProxyTypes): return False
+      return isinstance(o, internal.NativeObject)
+
+    tensors = [
+        o for o in gc.get_objects() if is_native_object(o)
+    ]
+    self.assertEmpty(tensors, "%d Tensors are still alive." % len(tensors))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testEagerMemoryUsageWithReset(self):
+    if memory_profiler is None:
+      self.skipTest("memory_profiler required to run this test")
+
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]])
+
+    def f():
+      self.evaluate(multi_device_iterator.get_next())
+      multi_device_iterator._eager_reset()
+
+    self.assertMemoryNotIncreasing(f, num_iters=50, max_increase_mb=250)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testEagerMemoryUsageWithRecreation(self):
+    if memory_profiler is None:
+      self.skipTest("memory_profiler required to run this test")
+
+    dataset = dataset_ops.Dataset.range(10)
+
+    def f():
+      multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+          dataset, [self._devices[1], self._devices[2]])
+      self.evaluate(multi_device_iterator.get_next())
+      del multi_device_iterator
+
+    # TODO(b/123316347): Reduce threshold once bug is fixed.
+    self.assertMemoryNotIncreasing(f, num_iters=50, max_increase_mb=250)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testFilter(self):
+
+    def get_dataset():
+
+      def fn(_):
+        return True
+
+      return dataset_ops.Dataset.range(0, 100).filter(fn)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="eager"))
+  def testFilterLegacy(self):
+
+    def get_dataset():
+
+      def fn(_):
+        return True
+
+      return dataset_ops.Dataset.range(0, 100).filter_with_legacy_function(fn)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testFlatMap(self):
+
+    def get_dataset():
+
+      def fn(x):
+        return dataset_ops.Dataset.from_tensors(x * x)
+
+      return dataset_ops.Dataset.range(0, 100).flat_map(fn)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testFromGenerator(self):
+
+    def get_dataset():
+
+      def fn():
+        return range(100)
+
+      return dataset_ops.Dataset.from_generator(fn, output_types=dtypes.float32)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 10])))
+  def testMap(self, num_parallel_calls):
+
+    def get_dataset():
+
+      def fn(x):
+        return x * x
+
+      return dataset_ops.Dataset.range(0, 100).map(
+          fn, num_parallel_calls=num_parallel_calls)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(
+      combinations.combine(
+          tf_api_version=1, mode="eager", num_parallel_calls=[None, 10]))
+  def testMapLegacy(self, num_parallel_calls):
+
+    def get_dataset():
+
+      def fn(x):
+        return x * x
+
+      return dataset_ops.Dataset.range(0, 100).map_with_legacy_function(
+          fn, num_parallel_calls=num_parallel_calls)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_parallel_calls=[None, 10])))
+  def testInterleave(self, num_parallel_calls):
+
+    def get_dataset():
+
+      def fn(x):
+        return dataset_ops.Dataset.from_tensors(x * x)
+
+      return dataset_ops.Dataset.range(0, 100).interleave(
+          fn, num_parallel_calls=num_parallel_calls, cycle_length=10)
+
+    self.assertNoMemoryLeak(get_dataset)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/multi_device_iterator_test.py

@@ -0,0 +1,476 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the `MultiDeviceIterator` and `OwnedMultiDeviceIterator` API."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import from_generator_op
+from tensorflow.python.data.ops import multi_device_iterator_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import cancellation
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import executor
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.platform import test
+
+
+cls_combination = combinations.combine(cls=[
+    combinations.NamedObject("MultiDeviceIterator",
+                             multi_device_iterator_ops.MultiDeviceIterator),
+    combinations.NamedObject("OwnedMultiDeviceIterator",
+                             multi_device_iterator_ops.OwnedMultiDeviceIterator)
+])
+
+
+class MultiDeviceIteratorCommonTest(test_base.DatasetTestBase,
+                                    parameterized.TestCase):
+  """Tests that are common to MultiDeviceIterator and OwnedMultiDeviceIterator."""
+
+  def setUp(self):
+    super().setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(3)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(), cls_combination))
+  def testCancelGetNextWithDevice(self, cls):
+    ping = data_flow_ops.FIFOQueue(capacity=2, dtypes=dtypes.int64)
+    pong = data_flow_ops.FIFOQueue(capacity=2, dtypes=dtypes.int64)
+
+    @def_function.function
+    def map_fn(v):
+      ball = ping.dequeue()
+      with ops.control_dependencies([pong.enqueue(ball)]):
+        return v + ping.dequeue()
+
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.map(map_fn)
+
+    # We need to set prefetch_buffer_size=0 so that we can cancel the
+    # MultiDeviceIteratorGetNextFromShardOp from eager. If
+    # prefetch_buffer_size>0, that op runs in the background threads of the
+    # prefetch and can only be cancelled by deleting the iterator.
+    multi_device_iterator = cls(
+        dataset, [self._devices[1], self._devices[2]], prefetch_buffer_size=0)
+
+    @def_function.function
+    def get_next_device1():
+      return multi_device_iterator.get_next(self._devices[1])
+
+    async_executor = executor.new_executor(enable_async=True)
+    with context.executor_scope(async_executor):
+      cancel_mgr = cancellation.CancellationManager()
+      cancel_mgr.get_cancelable_function(
+          get_next_device1.get_concrete_function())()
+    # Make sure we cancel in the middle of get_next.
+    ping.enqueue(0)
+    pong.dequeue()
+    cancel_mgr.start_cancel()
+    with self.assertRaises(errors.CancelledError):
+      async_executor.wait()
+    # Note that fetching from upstream iterator is not cancelled with the
+    # cancellation of get_next.
+    ping.enqueue(0)
+
+    # Cancelling a get_next on one device shouldn't cancel the
+    # multi_device_iterator and iterators on other devices.
+    ping.enqueue(0)
+    ping.enqueue(0)
+    self.assertEqual(1,
+                     multi_device_iterator.get_next(self._devices[2]).numpy())
+    # FIXME(b/209534797): Workaround an asan error caused by this test.
+    # Remove the dangling reference from tf.function to ensure queue objects
+    # are not freed before they are flushed.
+    import gc  # pylint: disable=g-import-not-at-top
+    del get_next_device1
+    gc.collect()
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(), cls_combination))
+  def testEmptyDataset(self, cls):
+    dataset = dataset_ops.Dataset.range(0)
+    multi_device_iterator = cls(
+        dataset, devices=[self._devices[1], self._devices[2]])
+    with self.assertRaises(errors.OutOfRangeError):
+      multi_device_iterator.get_next()
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(), cls_combination))
+  def testEmptyDeviceList(self, cls):
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Length for attr 'devices' of 0 must be at least minimum 1"):
+      cls(dataset, devices=[])
+
+
+class MultiDeviceIteratorTest(test_base.DatasetTestBase,
+                              parameterized.TestCase):
+
+  def setUp(self):
+    super(MultiDeviceIteratorTest, self).setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_inits=[0, 1, 42])))
+  def testInitOnly(self, num_inits):
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]])
+
+    for _ in range(num_inits):
+      self.evaluate(multi_device_iterator.initializer)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              max_buffer_size=[0, 1, 10], prefetch_buffer_size=[0, 1, 10])))
+  def testBasic(self, prefetch_buffer_size, max_buffer_size):
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]],
+        max_buffer_size=max_buffer_size,
+        prefetch_buffer_size=prefetch_buffer_size)
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 10, 2):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.assertEqual(i, self.evaluate(elem_on_1))
+      self.assertEqual(i + 1, self.evaluate(elem_on_2))
+    with self.assertRaises(errors.OutOfRangeError):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOneOnSameDevice(self):
+    dataset = dataset_ops.Dataset.range(12)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[0], self._devices[1], self._devices[2]])
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 12, 3):
+      elem_on_0, elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.assertEqual(i, self.evaluate(elem_on_0))
+      self.assertEqual(i + 1, self.evaluate(elem_on_1))
+      self.assertEqual(i + 2, self.evaluate(elem_on_2))
+    with self.assertRaises(errors.OutOfRangeError):
+      elem_on_0, elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.evaluate(elem_on_0)
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRepeatDevices(self):
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[1]])
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 10, 2):
+      elements = multi_device_iterator.get_next()
+      elem_on_1, elem_on_2 = elements
+      self.assertEqual(i, self.evaluate(elem_on_1))
+      self.assertEqual(i + 1, self.evaluate(elem_on_2))
+    with self.assertRaises(errors.OutOfRangeError):
+      elements = multi_device_iterator.get_next()
+      elem_on_1, elem_on_2 = elements
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNotFullyDivisible(self):
+    dataset = dataset_ops.Dataset.range(9)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]])
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 8, 2):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.assertEqual(i, self.evaluate(elem_on_1))
+      self.assertEqual(i + 1, self.evaluate(elem_on_2))
+    elem_on_1 = multi_device_iterator.get_next(self._devices[1])
+    self.assertEqual(8, self.evaluate(elem_on_1))
+    with self.assertRaises(errors.OutOfRangeError):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGetNextAsOptional(self):
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]])
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 10, 2):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next_as_optional()
+      has_elem_1, get_elem_1 = self.evaluate(
+          [elem_on_1.has_value(), elem_on_1.get_value()])
+      has_elem_2, get_elem_2 = self.evaluate(
+          [elem_on_2.has_value(), elem_on_2.get_value()])
+      self.assertTrue(has_elem_1)
+      self.assertEqual(i, get_elem_1)
+      self.assertTrue(has_elem_2)
+      self.assertEqual(i + 1, get_elem_2)
+    elem_on_1, elem_on_2 = multi_device_iterator.get_next_as_optional()
+    has_elem_1 = elem_on_1.has_value()
+    has_elem_2 = elem_on_2.has_value()
+    self.assertFalse(self.evaluate(has_elem_1))
+    self.assertFalse(self.evaluate(has_elem_2))
+    with self.assertRaises(errors.InvalidArgumentError):
+      elem_1 = elem_on_1.get_value()
+      self.evaluate(elem_1)
+    with self.assertRaises(errors.InvalidArgumentError):
+      elem_2 = elem_on_2.get_value()
+      self.evaluate(elem_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUneven(self):
+    dataset = dataset_ops.Dataset.range(10)
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]], max_buffer_size=4)
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 10, 2):
+      elem_on_1 = multi_device_iterator.get_next(self._devices[1])
+      self.assertEqual(i, self.evaluate(elem_on_1))
+    for i in range(0, 10, 2):
+      elem_on_2 = multi_device_iterator.get_next(self._devices[2])
+      self.assertEqual(i + 1, self.evaluate(elem_on_2))
+    with self.assertRaises(errors.OutOfRangeError):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testMultipleInitializationsGraph(self):
+    dataset1 = dataset_ops.Dataset.range(1000)
+    dataset2 = dataset_ops.Dataset.range(1000)
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]], prefetch_buffer_size=4)
+    elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+
+    for _ in range(5):
+      self.evaluate(multi_device_iterator.initializer)
+      self.assertEqual([(0, 0), (1, 1)], self.evaluate([elem_on_1, elem_on_2]))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testMultipleInitializationsEager(self):
+    dataset1 = dataset_ops.Dataset.range(1000)
+    dataset2 = dataset_ops.Dataset.range(1000)
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+
+    for _ in range(5):
+      multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+          dataset, [self._devices[1], self._devices[2]], prefetch_buffer_size=4)
+      self.evaluate(multi_device_iterator.initializer)
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.assertEqual([(0, 0), (1, 1)], self.evaluate([elem_on_1, elem_on_2]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptimization(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(testing.assert_next(["MemoryCacheImpl"]))
+    dataset = dataset.skip(0)  # this should be optimized away
+    dataset = dataset.cache()
+
+    options = options_lib.Options()
+    options.experimental_optimization.noop_elimination = True
+    dataset = dataset.with_options(options)
+
+    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]])
+
+    self.evaluate(multi_device_iterator.initializer)
+    for i in range(0, 10, 2):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.assertEqual(i, self.evaluate(elem_on_1))
+      self.assertEqual(i + 1, self.evaluate(elem_on_2))
+    with self.assertRaises(errors.OutOfRangeError):
+      elem_on_1, elem_on_2 = multi_device_iterator.get_next()
+      self.evaluate(elem_on_1)
+      self.evaluate(elem_on_2)
+
+
+class OwnedMultiDeviceIteratorTest(test_base.DatasetTestBase,
+                                   parameterized.TestCase):
+
+  def setUp(self):
+    super(OwnedMultiDeviceIteratorTest, self).setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(3)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(
+              max_buffer_size=[0, 1, 10], prefetch_buffer_size=[0, 1, 10])))
+  def testBasic(self, max_buffer_size, prefetch_buffer_size):
+    dataset = dataset_ops.Dataset.range(1000)
+
+    mdi = multi_device_iterator_ops.OwnedMultiDeviceIterator(
+        dataset, [self._devices[1], self._devices[2]],
+        max_buffer_size=max_buffer_size,
+        prefetch_buffer_size=prefetch_buffer_size)
+
+    for i, el in enumerate(mdi):
+      self.assertEqual([i * 2, i * 2 + 1], [el[0].numpy(), el[1].numpy()])
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testBasicFunction(self):
+    queue = data_flow_ops.FIFOQueue(10, dtypes.int64)
+
+    @def_function.function
+    def fn():
+      with ops.device(self._devices[0]):
+        dataset = dataset_ops.Dataset.range(10)
+      iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
+          dataset, [self._devices[1], self._devices[2]])
+      for _ in range(5):
+        el0, el1 = next(iterator)
+        queue.enqueue(el0)
+        queue.enqueue(el1)
+
+    fn()
+
+    for i in range(10):
+      self.assertEqual(queue.dequeue().numpy(), i)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testFunctionError(self):
+    # In this test we verify that a function that raises an error ends up
+    # properly deallocating the iterator resource.
+
+    queue = data_flow_ops.FIFOQueue(10, dtypes.int64)
+    queue.enqueue(0)
+
+    def init_fn(n):
+      return n
+
+    def next_fn(_):
+      ds = dataset_ops.Dataset.range(0)
+      return next(iter(ds))
+
+    def finalize_fn(n):
+      queue.enqueue(0)
+      return n
+
+    @def_function.function
+    def fn():
+      dataset = from_generator_op._GeneratorDataset(
+          1,
+          init_fn,
+          next_fn,
+          finalize_fn,
+          output_signature=tensor_spec.TensorSpec([], dtypes.int64))
+      iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
+          dataset, [self._devices[1], self._devices[2]])
+      next(iterator)
+
+    with self.assertRaises(errors.OutOfRangeError):
+      fn()
+
+    self.assertEqual(queue.size().numpy(), 2)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testMultipleInitializations(self):
+    dataset = dataset_ops.Dataset.range(1000)
+
+    for _ in range(5):
+      multi_device_iterator = (
+          multi_device_iterator_ops.OwnedMultiDeviceIterator(
+              dataset, [self._devices[1], self._devices[2]]))
+      for i, el in enumerate(multi_device_iterator):
+        self.assertEqual([i * 2, i * 2 + 1], [el[0].numpy(), el[1].numpy()])
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testLimitedRetracing(self):
+    trace_count = [0]
+
+    @def_function.function
+    def f(iterator):
+      trace_count[0] += 1
+      counter = np.int64(0)
+      for _ in range(5):
+        elem = next(iterator)
+        counter += elem[0]
+        counter += elem[1]
+      return counter
+
+    dataset = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(20)
+
+    for _ in range(10):
+      multi_device_iterator = (
+          multi_device_iterator_ops.OwnedMultiDeviceIterator(
+              dataset, [self._devices[1], self._devices[2]]))
+      self.assertEqual(self.evaluate(f(multi_device_iterator)), 45)
+      multi_device_iterator2 = (
+          multi_device_iterator_ops.OwnedMultiDeviceIterator(
+              dataset2, [self._devices[1], self._devices[2]]))
+      self.assertEqual(self.evaluate(f(multi_device_iterator2)), 45)
+      self.assertEqual(trace_count[0], 1)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testMissingDevices(self):
+    dataset = dataset_ops.Dataset.range(1000)
+    with self.assertRaisesRegex(ValueError, "`devices` must be provided."):
+      multi_device_iterator_ops.OwnedMultiDeviceIterator(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testMissingInput(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        "When `dataset` is not provided, both `components` and `element_spec` "
+        "must be specified."):
+      multi_device_iterator_ops.OwnedMultiDeviceIterator(
+          dataset=None, devices=[self._devices[1], self._devices[2]])
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testExtraElementSpecInput(self):
+    dataset = dataset_ops.Dataset.range(1000)
+    with self.assertRaisesRegex(
+        ValueError,
+        "When `dataset` is provided, `element_spec` and `components` must "
+        "not be specified."):
+      multi_device_iterator_ops.OwnedMultiDeviceIterator(
+          dataset, devices=[self._devices[1], self._devices[2]],
+          element_spec=dataset.element_spec)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testGraphMode(self):
+    dataset = dataset_ops.Dataset.range(1000)
+    with self.assertRaisesRegex(
+        RuntimeError,
+        "OwnedMultiDeviceIterator is only supported inside of tf.function or "
+        "when eager execution is enabled."):
+      multi_device_iterator_ops.OwnedMultiDeviceIterator(
+          dataset, devices=[self._devices[1], self._devices[2]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/optional_test.py

@@ -0,0 +1,460 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Optional`."""
+import functools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.data.ops import optional_ops
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _optional_spec_test_combinations():
+  # pylint: disable=g-long-lambda
+  cases = [
+      ("Dense", lambda: constant_op.constant(37.0),
+       tensor_spec.TensorSpec([], dtypes.float32)),
+      ("Sparse", lambda: sparse_tensor.SparseTensor(
+          indices=[[0, 1]],
+          values=constant_op.constant([0], dtype=dtypes.int32),
+          dense_shape=[10, 10]),
+       sparse_tensor.SparseTensorSpec([10, 10], dtypes.int32)),
+      ("Nest", lambda: {
+          "a": constant_op.constant(37.0),
+          "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar"))
+      }, {
+          "a":
+              tensor_spec.TensorSpec([], dtypes.float32),
+          "b": (
+              tensor_spec.TensorSpec([1], dtypes.string),
+              tensor_spec.TensorSpec([], dtypes.string),
+          )
+      }),
+      ("Optional", lambda: optional_ops.Optional.from_value(37.0),
+       optional_ops.OptionalSpec(tensor_spec.TensorSpec([], dtypes.float32))),
+  ]
+
+  def reduce_fn(x, y):
+    name, value_fn, expected_structure = y
+    return x + combinations.combine(
+        tf_value_fn=combinations.NamedObject(name, value_fn),
+        expected_value_structure=expected_structure)
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _get_next_as_optional_test_combinations():
+  # pylint: disable=g-long-lambda
+  cases = [
+      ("Dense", np.array([1, 2, 3], dtype=np.int32),
+       lambda: constant_op.constant([4, 5, 6], dtype=dtypes.int32), True),
+      ("Sparse",
+       sparse_tensor.SparseTensorValue(
+           indices=[[0, 0], [1, 1]],
+           values=np.array([-1., 1.], dtype=np.float32),
+           dense_shape=[2, 2]),
+       lambda: sparse_tensor.SparseTensor(
+           indices=[[0, 1], [1, 0]], values=[37.0, 42.0], dense_shape=[2, 2]),
+       False),
+      ("Nest", {
+          "a":
+              np.array([1, 2, 3], dtype=np.int32),
+          "b":
+              sparse_tensor.SparseTensorValue(
+                  indices=[[0, 0], [1, 1]],
+                  values=np.array([-1., 1.], dtype=np.float32),
+                  dense_shape=[2, 2])
+      }, lambda: {
+          "a":
+              constant_op.constant([4, 5, 6], dtype=dtypes.int32),
+          "b":
+              sparse_tensor.SparseTensor(
+                  indices=[[0, 1], [1, 0]],
+                  values=[37.0, 42.0],
+                  dense_shape=[2, 2])
+      }, False),
+  ]
+
+  def reduce_fn(x, y):
+    name, value, value_fn, gpu_compatible = y
+    return x + combinations.combine(
+        np_value=value,
+        tf_value_fn=combinations.NamedObject(name, value_fn),
+        gpu_compatible=gpu_compatible)
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromValue(self):
+    opt = optional_ops.Optional.from_value(constant_op.constant(37.0))
+    self.assertTrue(self.evaluate(opt.has_value()))
+    self.assertEqual(37.0, self.evaluate(opt.get_value()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromStructuredValue(self):
+    opt = optional_ops.Optional.from_value({
+        "a": constant_op.constant(37.0),
+        "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar"))
+    })
+    self.assertTrue(self.evaluate(opt.has_value()))
+    self.assertEqual({
+        "a": 37.0,
+        "b": ([b"Foo"], b"Bar")
+    }, self.evaluate(opt.get_value()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromSparseTensor(self):
+    st_0 = sparse_tensor.SparseTensorValue(
+        indices=np.array([[0]]),
+        values=np.array([0], dtype=np.int64),
+        dense_shape=np.array([1]))
+    st_1 = sparse_tensor.SparseTensorValue(
+        indices=np.array([[0, 0], [1, 1]]),
+        values=np.array([-1., 1.], dtype=np.float32),
+        dense_shape=np.array([2, 2]))
+    opt = optional_ops.Optional.from_value((st_0, st_1))
+    self.assertTrue(self.evaluate(opt.has_value()))
+    val_0, val_1 = opt.get_value()
+    for expected, actual in [(st_0, val_0), (st_1, val_1)]:
+      self.assertAllEqual(expected.indices, self.evaluate(actual.indices))
+      self.assertAllEqual(expected.values, self.evaluate(actual.values))
+      self.assertAllEqual(expected.dense_shape,
+                          self.evaluate(actual.dense_shape))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromNone(self):
+    value_structure = tensor_spec.TensorSpec([], dtypes.float32)
+    opt = optional_ops.Optional.empty(value_structure)
+    self.assertTrue(opt.element_spec.is_compatible_with(value_structure))
+    self.assertFalse(
+        opt.element_spec.is_compatible_with(
+            tensor_spec.TensorSpec([1], dtypes.float32)))
+    self.assertFalse(
+        opt.element_spec.is_compatible_with(
+            tensor_spec.TensorSpec([], dtypes.int32)))
+    self.assertFalse(self.evaluate(opt.has_value()))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(opt.get_value())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAddN(self):
+    devices = ["/cpu:0"]
+    if test_util.is_gpu_available():
+      devices.append("/gpu:0")
+    for device in devices:
+      with ops.device(device):
+        # With value
+        opt1 = optional_ops.Optional.from_value((1.0, 2.0))
+        opt2 = optional_ops.Optional.from_value((3.0, 4.0))
+
+        add_tensor = math_ops.add_n(
+            [opt1._variant_tensor, opt2._variant_tensor])
+        add_opt = optional_ops._OptionalImpl(add_tensor, opt1.element_spec)
+        self.assertAllEqual(self.evaluate(add_opt.get_value()), (4.0, 6.0))
+
+        # Without value
+        opt_none1 = optional_ops.Optional.empty(opt1.element_spec)
+        opt_none2 = optional_ops.Optional.empty(opt2.element_spec)
+        add_tensor = math_ops.add_n(
+            [opt_none1._variant_tensor, opt_none2._variant_tensor])
+        add_opt = optional_ops._OptionalImpl(add_tensor, opt_none1.element_spec)
+        self.assertFalse(self.evaluate(add_opt.has_value()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedAddN(self):
+    devices = ["/cpu:0"]
+    if test_util.is_gpu_available():
+      devices.append("/gpu:0")
+    for device in devices:
+      with ops.device(device):
+        opt1 = optional_ops.Optional.from_value([1, 2.0])
+        opt2 = optional_ops.Optional.from_value([3, 4.0])
+        opt3 = optional_ops.Optional.from_value((5.0, opt1._variant_tensor))
+        opt4 = optional_ops.Optional.from_value((6.0, opt2._variant_tensor))
+
+        add_tensor = math_ops.add_n(
+            [opt3._variant_tensor, opt4._variant_tensor])
+        add_opt = optional_ops._OptionalImpl(add_tensor, opt3.element_spec)
+        self.assertEqual(self.evaluate(add_opt.get_value()[0]), 11.0)
+
+        inner_add_opt = optional_ops._OptionalImpl(add_opt.get_value()[1],
+                                                   opt1.element_spec)
+        self.assertAllEqual(inner_add_opt.get_value(), [4, 6.0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZerosLike(self):
+    devices = ["/cpu:0"]
+    if test_util.is_gpu_available():
+      devices.append("/gpu:0")
+    for device in devices:
+      with ops.device(device):
+        # With value
+        opt = optional_ops.Optional.from_value((1.0, 2.0))
+        zeros_tensor = array_ops.zeros_like(opt._variant_tensor)
+        zeros_opt = optional_ops._OptionalImpl(zeros_tensor, opt.element_spec)
+        self.assertAllEqual(self.evaluate(zeros_opt.get_value()), (0.0, 0.0))
+
+        # Without value
+        opt_none = optional_ops.Optional.empty(opt.element_spec)
+        zeros_tensor = array_ops.zeros_like(opt_none._variant_tensor)
+        zeros_opt = optional_ops._OptionalImpl(zeros_tensor,
+                                               opt_none.element_spec)
+        self.assertFalse(self.evaluate(zeros_opt.has_value()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedZerosLike(self):
+    devices = ["/cpu:0"]
+    if test_util.is_gpu_available():
+      devices.append("/gpu:0")
+    for device in devices:
+      with ops.device(device):
+        opt1 = optional_ops.Optional.from_value(1.0)
+        opt2 = optional_ops.Optional.from_value(opt1._variant_tensor)
+
+        zeros_tensor = array_ops.zeros_like(opt2._variant_tensor)
+        zeros_opt = optional_ops._OptionalImpl(zeros_tensor, opt2.element_spec)
+        inner_zeros_opt = optional_ops._OptionalImpl(zeros_opt.get_value(),
+                                                     opt1.element_spec)
+        self.assertEqual(self.evaluate(inner_zeros_opt.get_value()), 0.0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCopyToGPU(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    with ops.device("/cpu:0"):
+      optional_with_value = optional_ops.Optional.from_value(
+          (constant_op.constant(37.0), constant_op.constant("Foo"),
+           constant_op.constant(42)))
+      optional_none = optional_ops.Optional.empty(
+          tensor_spec.TensorSpec([], dtypes.float32))
+
+    with ops.device("/gpu:0"):
+      gpu_optional_with_value = optional_ops._OptionalImpl(
+          array_ops.identity(optional_with_value._variant_tensor),
+          optional_with_value.element_spec)
+      gpu_optional_none = optional_ops._OptionalImpl(
+          array_ops.identity(optional_none._variant_tensor),
+          optional_none.element_spec)
+
+      gpu_optional_with_value_has_value = gpu_optional_with_value.has_value()
+      gpu_optional_with_value_values = gpu_optional_with_value.get_value()
+
+      gpu_optional_none_has_value = gpu_optional_none.has_value()
+
+    self.assertTrue(self.evaluate(gpu_optional_with_value_has_value))
+    self.assertEqual((37.0, b"Foo", 42),
+                     self.evaluate(gpu_optional_with_value_values))
+    self.assertFalse(self.evaluate(gpu_optional_none_has_value))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedCopyToGPU(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    with ops.device("/cpu:0"):
+      optional_with_value = optional_ops.Optional.from_value(
+          (constant_op.constant(37.0), constant_op.constant("Foo"),
+           constant_op.constant(42)))
+      optional_none = optional_ops.Optional.empty(
+          tensor_spec.TensorSpec([], dtypes.float32))
+      nested_optional = optional_ops.Optional.from_value(
+          (optional_with_value._variant_tensor, optional_none._variant_tensor,
+           1.0))
+
+    with ops.device("/gpu:0"):
+      gpu_nested_optional = optional_ops._OptionalImpl(
+          array_ops.identity(nested_optional._variant_tensor),
+          nested_optional.element_spec)
+
+      gpu_nested_optional_has_value = gpu_nested_optional.has_value()
+      gpu_nested_optional_values = gpu_nested_optional.get_value()
+
+    self.assertTrue(self.evaluate(gpu_nested_optional_has_value))
+
+    inner_with_value = optional_ops._OptionalImpl(
+        gpu_nested_optional_values[0], optional_with_value.element_spec)
+
+    inner_none = optional_ops._OptionalImpl(gpu_nested_optional_values[1],
+                                            optional_none.element_spec)
+
+    self.assertEqual((37.0, b"Foo", 42),
+                     self.evaluate(inner_with_value.get_value()))
+    self.assertFalse(self.evaluate(inner_none.has_value()))
+    self.assertEqual(1.0, self.evaluate(gpu_nested_optional_values[2]))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _optional_spec_test_combinations()))
+  def testOptionalSpec(self, tf_value_fn, expected_value_structure):
+    tf_value = tf_value_fn()
+    opt = optional_ops.Optional.from_value(tf_value)
+
+    self.assertTrue(
+        structure.are_compatible(opt.element_spec, expected_value_structure))
+
+    opt_structure = structure.type_spec_from_value(opt)
+    self.assertIsInstance(opt_structure, optional_ops.OptionalSpec)
+    self.assertTrue(structure.are_compatible(opt_structure, opt_structure))
+    self.assertTrue(
+        structure.are_compatible(opt_structure._element_spec,
+                                 expected_value_structure))
+    self.assertEqual([dtypes.variant],
+                     structure.get_flat_tensor_types(opt_structure))
+    self.assertEqual([tensor_shape.TensorShape([])],
+                     structure.get_flat_tensor_shapes(opt_structure))
+
+    # All OptionalSpec objects are not compatible with a non-optional
+    # value.
+    non_optional_structure = structure.type_spec_from_value(
+        constant_op.constant(42.0))
+    self.assertFalse(opt_structure.is_compatible_with(non_optional_structure))
+
+    # Assert that the optional survives a round-trip via _from_tensor_list()
+    # and _to_tensor_list().
+    round_trip_opt = opt_structure._from_tensor_list(
+        opt_structure._to_tensor_list(opt))
+    if isinstance(tf_value, optional_ops.Optional):
+      self.assertValuesEqual(
+          self.evaluate(tf_value.get_value()),
+          self.evaluate(round_trip_opt.get_value().get_value()))
+    else:
+      self.assertValuesEqual(
+          self.evaluate(tf_value), self.evaluate(round_trip_opt.get_value()))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _get_next_as_optional_test_combinations()))
+  def testIteratorGetNextAsOptional(self, np_value, tf_value_fn,
+                                    gpu_compatible):
+    if not gpu_compatible and test.is_gpu_available():
+      self.skipTest("Test case not yet supported on GPU.")
+    ds = dataset_ops.Dataset.from_tensors(np_value).repeat(3)
+
+    if context.executing_eagerly():
+      iterator = dataset_ops.make_one_shot_iterator(ds)
+      # For each element of the dataset, assert that the optional evaluates to
+      # the expected value.
+      for _ in range(3):
+        next_elem = iterator_ops.get_next_as_optional(iterator)
+        self.assertIsInstance(next_elem, optional_ops.Optional)
+        self.assertTrue(
+            structure.are_compatible(
+                next_elem.element_spec,
+                structure.type_spec_from_value(tf_value_fn())))
+        self.assertTrue(next_elem.has_value())
+        self.assertValuesEqual(np_value, next_elem.get_value())
+      # After exhausting the iterator, `next_elem.has_value()` will evaluate to
+      # false, and attempting to get the value will fail.
+      for _ in range(2):
+        next_elem = iterator_ops.get_next_as_optional(iterator)
+        self.assertFalse(self.evaluate(next_elem.has_value()))
+        with self.assertRaises(errors.InvalidArgumentError):
+          self.evaluate(next_elem.get_value())
+    else:
+      iterator = dataset_ops.make_initializable_iterator(ds)
+      next_elem = iterator_ops.get_next_as_optional(iterator)
+      self.assertIsInstance(next_elem, optional_ops.Optional)
+      self.assertTrue(
+          structure.are_compatible(
+              next_elem.element_spec,
+              structure.type_spec_from_value(tf_value_fn())))
+      # Before initializing the iterator, evaluating the optional fails with
+      # a FailedPreconditionError. This is only relevant in graph mode.
+      elem_has_value_t = next_elem.has_value()
+      elem_value_t = next_elem.get_value()
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(elem_has_value_t)
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(elem_value_t)
+      # Now we initialize the iterator.
+      self.evaluate(iterator.initializer)
+      # For each element of the dataset, assert that the optional evaluates to
+      # the expected value.
+      for _ in range(3):
+        elem_has_value, elem_value = self.evaluate(
+            [elem_has_value_t, elem_value_t])
+        self.assertTrue(elem_has_value)
+        self.assertValuesEqual(np_value, elem_value)
+
+      # After exhausting the iterator, `next_elem.has_value()` will evaluate to
+      # false, and attempting to get the value will fail.
+      for _ in range(2):
+        self.assertFalse(self.evaluate(elem_has_value_t))
+        with self.assertRaises(errors.InvalidArgumentError):
+          self.evaluate(elem_value_t)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFunctionBoundaries(self):
+
+    @def_function.function
+    def get_optional():
+      x = constant_op.constant(1.0)
+      opt = optional_ops.Optional.from_value(x)
+      # TODO(skyewm): support returning Optionals from functions?
+      return opt._variant_tensor
+
+    # TODO(skyewm): support Optional arguments?
+    @def_function.function
+    def consume_optional(opt_tensor):
+      value_structure = tensor_spec.TensorSpec([], dtypes.float32)
+      opt = optional_ops._OptionalImpl(opt_tensor, value_structure)
+      return opt.get_value()
+
+    opt_tensor = get_optional()
+    val = consume_optional(opt_tensor)
+    self.assertEqual(self.evaluate(val), 1.0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLimitedRetracing(self):
+    trace_count = [0]
+
+    @def_function.function
+    def f(opt):
+      trace_count[0] += 1
+      return opt.get_value()
+
+    opt1 = optional_ops.Optional.from_value(constant_op.constant(37.0))
+    opt2 = optional_ops.Optional.from_value(constant_op.constant(42.0))
+
+    for _ in range(10):
+      self.assertEqual(self.evaluate(f(opt1)), 37.0)
+      self.assertEqual(self.evaluate(f(opt2)), 42.0)
+      self.assertEqual(trace_count[0], 1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/options_test.py

@@ -0,0 +1,311 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Options`."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.framework import dataset_options_pb2
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+class OptionsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _get_options(self, dataset):
+    if context.executing_eagerly():
+      return dataset.options()
+    return dataset_ops.Dataset._options_tensor_to_options(
+        self.evaluate(dataset._options()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsDefault(self):
+    ds = dataset_ops.Dataset.range(0)
+    self.assertEqual(options_lib.Options(), ds.options())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsOnce(self):
+    options = options_lib.Options()
+    ds = dataset_ops.Dataset.range(0).with_options(options).cache()
+    self.assertEqual(options, ds.options())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsTwiceSame(self):
+    options = options_lib.Options()
+    options.autotune.enabled = True
+    ds = dataset_ops.Dataset.range(0).with_options(options).with_options(
+        options)
+    self.assertEqual(options, self._get_options(ds))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsTwiceDifferentOptions(self):
+    options1 = options_lib.Options()
+    options1.autotune.enabled = True
+    options2 = options_lib.Options()
+    options2.deterministic = False
+    ds = dataset_ops.Dataset.range(0)
+    ds = ds.with_options(options1)
+    ds = ds.with_options(options2)
+    options = self._get_options(ds)
+    self.assertTrue(options.autotune.enabled)
+    # Explicitly check that flag is False since assertFalse allows None
+    self.assertIs(options.deterministic, False)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsTwiceSameOption(self):
+    options1 = options_lib.Options()
+    options1.autotune.enabled = False
+    options2 = options_lib.Options()
+    options2.autotune.enabled = True
+    ds = dataset_ops.Dataset.range(0)
+    ds = ds.with_options(options1)
+    ds = ds.with_options(options2)
+    self.assertTrue(self._get_options(ds).autotune.enabled)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsTwiceSameOptionWithMap(self):
+    options1 = options_lib.Options()
+    options1.framework_type = ["seqio"]
+    options2 = options_lib.Options()
+    options2.framework_type = ["tfgrain"]
+    ds = dataset_ops.Dataset.range(5)
+    ds = ds.with_options(options1)
+    ds = ds.map(lambda x: x + 1)
+    ds = ds.with_options(options2)
+    self.assertDatasetProduces(ds, [1, 2, 3, 4, 5])
+    self.assertLen(self._get_options(ds).framework_type, 2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsMergeOptionsFromMultipleInputs(self):
+    options1 = options_lib.Options()
+    options1.autotune.enabled = True
+    options2 = options_lib.Options()
+    options2.deterministic = True
+    ds1 = dataset_ops.Dataset.range(0).with_options(options1)
+    ds2 = dataset_ops.Dataset.range(0).with_options(options2)
+    ds = dataset_ops.Dataset.zip((ds1, ds2))
+    options = self._get_options(ds)
+    self.assertTrue(options.autotune.enabled)
+    self.assertTrue(options.deterministic)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsHaveDefaults(self):
+    options1 = options_lib.Options()
+    options2 = options_lib.Options()
+    self.assertIsNot(options1.experimental_optimization,
+                     options2.experimental_optimization)
+    self.assertIsNot(options1.threading, options2.threading)
+    self.assertEqual(options1.experimental_optimization,
+                     options_lib.OptimizationOptions())
+    self.assertEqual(options1.threading, options_lib.ThreadingOptions())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMutatingOptionsRaiseValueError(self):
+    ds = dataset_ops.Dataset.range(0)
+    options1 = options_lib.Options()
+    options1.experimental_slack = True
+    options2 = options_lib.Options()
+    options2.autotune.enabled = True
+    ds = ds.with_options(options1)
+    ds = ds.map(lambda x: 2 * x)
+    ds = ds.with_options(options2)
+    dataset_options = ds.options()
+    with self.assertRaises(ValueError):
+      dataset_options.deterministic = True
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testNestedDataset(self):
+    ds = dataset_ops.Dataset.from_tensors(0)
+    result = ds
+
+    for _ in range(99):
+      result = result.concatenate(ds)
+    self.assertDatasetProduces(result, [0]*100)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsProtoRoundTrip(self):
+    options = options_lib.Options()
+    options.autotune.enabled = True
+    options.autotune.cpu_budget = 10
+    options.autotune.ram_budget = 20
+    options.deterministic = True
+    options.experimental_external_state_policy = (
+        options_lib.ExternalStatePolicy.FAIL)
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.DATA)
+    options.experimental_distribute.num_devices = 1000
+    options.experimental_optimization.apply_default_optimizations = True
+    options.experimental_optimization.filter_fusion = True
+    options.experimental_optimization.filter_parallelization = True
+    options.experimental_optimization.inject_prefetch = False
+    options.experimental_optimization.map_and_batch_fusion = True
+    options.experimental_optimization.map_and_filter_fusion = True
+    options.experimental_optimization.map_fusion = True
+    options.experimental_optimization.map_parallelization = True
+    options.experimental_optimization.noop_elimination = True
+    options.experimental_optimization.parallel_batch = True
+    options.experimental_optimization.shuffle_and_repeat_fusion = True
+    options.experimental_optimization.seq_interleave_prefetch = True
+    options.experimental_warm_start = True
+    options.experimental_slack = True
+    options.dataset_name = "test_name"
+    options.framework_type = ["TFDS", "TfGrain"]
+    options.threading.max_intra_op_parallelism = 30
+    options.threading.private_threadpool_size = 40
+    pb = options._to_proto()
+    result = options_lib.Options()
+    result._from_proto(pb)
+    self.assertEqual(options.framework_type, result.framework_type)
+    self.assertEqual(options, result)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsProtoDefaultValuesRoundTrip(self):
+    options = options_lib.Options()
+    pb = options._to_proto()
+    result = options_lib.Options()
+    result._from_proto(pb)
+    self.assertEqual(options, result)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testProtoOptionsDefaultValuesRoundTrip(self):
+    pb = dataset_options_pb2.Options()
+    options = options_lib.Options()
+    options._from_proto(pb)
+    result = options._to_proto()
+    expected_pb = dataset_options_pb2.Options()
+    expected_pb.autotune_options.CopyFrom(dataset_options_pb2.AutotuneOptions())
+    expected_pb.distribute_options.CopyFrom(
+        dataset_options_pb2.DistributeOptions())
+    expected_pb.optimization_options.CopyFrom(
+        dataset_options_pb2.OptimizationOptions())
+    expected_pb.warm_start = True
+    expected_pb.service_options.CopyFrom(
+        dataset_options_pb2.ServiceOptions())
+    expected_pb.threading_options.CopyFrom(
+        dataset_options_pb2.ThreadingOptions())
+    self.assertProtoEquals(expected_pb, result)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testThreadingOptionsBackwardCompatibility(self):
+    opts = options_lib.Options()
+    opts.threading.max_intra_op_parallelism = 20
+    self.assertEqual(opts.experimental_threading.max_intra_op_parallelism, 20)
+    opts.experimental_threading.private_threadpool_size = 80
+    self.assertEqual(opts.threading.private_threadpool_size, 80)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExperimentalThreadingOptionsOverride(self):
+    options = options_lib.Options()
+    self.assertEqual(options.threading, options.experimental_threading)
+    options.threading.max_intra_op_parallelism = 20
+    options.experimental_threading.max_intra_op_parallelism = 40
+    pb = options._to_proto()
+    result = options_lib.Options()
+    result._from_proto(pb)
+    self.assertEqual(result.experimental_threading.max_intra_op_parallelism,
+                     result.threading.max_intra_op_parallelism)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExperimentalDeterministicOverride(self):
+    options = options_lib.Options()
+    self.assertEqual(options.deterministic, options.experimental_deterministic)
+    options.experimental_deterministic = False
+    pb = options._to_proto()
+    result = options_lib.Options()
+    result._from_proto(pb)
+    self.assertFalse(result.deterministic)
+    self.assertEqual(result.deterministic, result.experimental_deterministic)
+    result.experimental_deterministic = True
+    self.assertTrue(result.deterministic)
+    self.assertEqual(result.deterministic, result.experimental_deterministic)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPersistenceOptionsSetOutsideFunction(self):
+
+    @def_function.function
+    def fn(dataset):
+      dataset = dataset.map(lambda x: 10 * x)
+      return dataset
+
+    dataset = dataset_ops.Dataset.range(5)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    dataset = dataset.with_options(options)
+    dataset = fn(dataset)
+    result = dataset_ops.Dataset._options_tensor_to_options(
+        self.evaluate(dataset._options()))
+    self.assertTrue(result.experimental_slack)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPersistenceOptionsSetInsideFunction(self):
+
+    @def_function.function
+    def fn(dataset):
+      options = options_lib.Options()
+      options.experimental_slack = True
+      dataset = dataset.with_options(options)
+      dataset = dataset.map(lambda x: 10 * x)
+      return dataset
+
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = fn(dataset)
+    result = dataset_ops.Dataset._options_tensor_to_options(
+        self.evaluate(dataset._options()))
+    self.assertTrue(result.experimental_slack)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptionsPersistenceGraphRoundTrip(self):
+    dataset = dataset_ops.Dataset.range(5)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    dataset = self.graphRoundTrip(dataset)
+    result = self._get_options(dataset)
+    self.assertTrue(result.experimental_slack)
+    # Explicitly check that flag is False since assertFalse allows None
+    self.assertIs(
+        result.experimental_optimization.apply_default_optimizations, False)
+
+  @combinations.generate(combinations.times(
+      test_base.default_test_combinations(),
+      combinations.combine(map_parallelization=[True, False])))
+  def testOptionsGraphRoundTripOptimization(self, map_parallelization):
+    dataset = dataset_ops.Dataset.range(6)
+    options = options_lib.Options()
+    options.experimental_optimization.map_parallelization = (
+        map_parallelization)
+    dataset = dataset.with_options(options)
+    dataset = self.graphRoundTrip(dataset)
+    expected = "ParallelMap" if map_parallelization else "Map"
+    dataset = dataset.apply(testing.assert_next([expected]))
+    dataset = dataset.map(lambda x: x*x)
+    self.assertDatasetProduces(dataset, expected_output=[0, 1, 4, 9, 16, 25])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42)
+    options = options_lib.Options()
+    dataset = dataset.with_options(options, name="options")
+    self.assertDatasetProduces(dataset, [42])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/padded_batch_test.py

@@ -0,0 +1,412 @@
+# -*- coding: utf-8 -*-
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.padded_batch()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              count=[32, 34],
+              padded_shapes=[[None], [25]],
+              drop_remainder=[True, False])))
+  def testPaddedBatchDataset(self, count, padded_shapes, drop_remainder):
+    seq_lens = np.random.randint(20, size=(count,)).astype(np.int32)
+    batch_size = 4
+    dataset = dataset_ops.Dataset.from_tensor_slices(seq_lens).map(
+        lambda x: array_ops.fill([x], x)).padded_batch(
+            batch_size=batch_size,
+            drop_remainder=drop_remainder,
+            padded_shapes=padded_shapes)
+
+    num_full_batches = len(seq_lens) // batch_size
+    get_next = self.getNext(dataset)
+    for i in range(num_full_batches):
+      result = self.evaluate(get_next())
+      padded_len = padded_shapes[0]
+      if padded_len is None or padded_len == -1:
+        padded_len = np.max(result) if result.size > 0 else 0
+      self.assertEqual((batch_size, padded_len), result.shape)
+      for j in range(batch_size):
+        seq_len = seq_lens[(i * batch_size) + j]
+        self.assertAllEqual(result[j, :seq_len], [seq_len] * seq_len)
+        self.assertAllEqual(result[j, seq_len:], [0] * (padded_len - seq_len))
+
+    if not drop_remainder and len(seq_lens) % batch_size > 0:
+      result = self.evaluate(get_next())
+      padded_len = padded_shapes[0]
+      if padded_len is None or padded_len == -1:
+        padded_len = np.max(result) if result.size > 0 else 0
+      self.assertEqual((len(seq_lens) % batch_size, padded_len), result.shape)
+      for j in range(len(seq_lens) % batch_size):
+        seq_len = seq_lens[num_full_batches * batch_size + j]
+        self.assertAllEqual(result[j, :seq_len], [seq_len] * seq_len)
+        self.assertAllEqual(result[j, seq_len:], [0] * (padded_len - seq_len))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShortPadding(self):
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(
+            [6, 5, 5, 5, 5]).map(lambda x: array_ops.fill([x], x)).padded_batch(
+                batch_size=4, padded_shapes=[5]))
+    self.assertDatasetProduces(
+        dataset, expected_error=(errors.DataLossError, ''))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchEmptyTensors(self):
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(
+            [0, 0, 0, 0]).map(lambda x: array_ops.fill([x], x)).padded_batch(
+                batch_size=4, padded_shapes=[-1]))
+    self.assertDatasetProduces(dataset, expected_output=[[[], [], [], []]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchWithDifferetElementTypes(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        ([0, 1, 2, 3], ['a', 'b', 'c', 'd']))
+    dataset = dataset.padded_batch(2)
+    self.assertDatasetProduces(dataset, [([0, 1], ['a', 'b']),
+                                         ([2, 3], ['c', 'd'])])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchWithEmptyTuple(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(([0, 1, 2, 3], ()))
+    dataset = dataset.padded_batch(2)
+    self.assertDatasetProduces(dataset, [([0, 1], ()), ([2, 3], ())])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchWithNoneElement(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(([0, 1, 2, 3], None))
+    with self.assertRaises(TypeError):
+      dataset = dataset.padded_batch(2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDefaultPaddedShapes(self):
+
+    def fill(x):
+      return array_ops.fill([x], x)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(
+            [1, 2, 3, 4]).map(fill).padded_batch(batch_size=2))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[[[1, 0], [2, 2]], [[3, 3, 3, 0], [4, 4, 4, 4]]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedDefaultPaddedShapes(self):
+
+    def fill_tuple(x):
+      return (x, array_ops.fill([x], x))
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(
+            [1, 2, 3, 4]).map(fill_tuple).padded_batch(batch_size=2))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[([1, 2], [[1, 0], [2, 2]]),
+                         ([3, 4], [[3, 3, 3, 0], [4, 4, 4, 4]])])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              padding_values=[(-1, '<end>', {'structure': ''}),
+                              (-1, '<end>', None)])))
+  def testPaddedBatchDatasetNonDefaultPadding(self, padding_values):
+
+    def fill_tuple(x):
+      filled = array_ops.fill([x], x)
+      return (filled, string_ops.as_string(filled), {
+          'structure': string_ops.as_string(filled)
+      })
+
+    random_seq_lens = np.random.randint(20, size=(32,)).astype(np.int32)
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(random_seq_lens).map(fill_tuple)
+        .padded_batch(
+            4, padded_shapes=([-1], [-1], {'structure': [-1]}),
+            padding_values=padding_values))
+
+    get_next = self.getNext(dataset)
+    for i in range(8):
+      result = self.evaluate(get_next())
+      padded_len = np.max(result[0])
+      self.assertEqual((4, padded_len), result[0].shape)
+      self.assertEqual((4, padded_len), result[1].shape)
+      self.assertEqual((4, padded_len), result[2]['structure'].shape)
+      for j in range(4):
+        seq_len = random_seq_lens[(i * 4) + j]
+        self.assertAllEqual(result[0][j, :seq_len], [seq_len] * seq_len)
+        self.assertAllEqual(result[0][j, seq_len:],
+                            [-1] * (padded_len - seq_len))
+        self.assertAllEqual(result[1][j, :seq_len],
+                            [compat.as_bytes(str(seq_len))] * seq_len)
+        self.assertAllEqual(result[1][j, seq_len:],
+                            [b'<end>'] * (padded_len - seq_len))
+        self.assertAllEqual(result[2]['structure'][j, :seq_len],
+                            [compat.as_bytes(str(seq_len))] * seq_len)
+        self.assertAllEqual(result[2]['structure'][j, seq_len:],
+                            [b''] * (padded_len - seq_len))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchDatasetUnicode(self):
+    # See GitHub issue 16149
+    def generator():
+      data = [[u'Простой', u'тест', u'юникода'],
+              [u'никогда', u'не', u'бывает', u'простым']]
+
+      for seq in data:
+        yield seq, [0, 1, 2, 3]
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, (dtypes.string, dtypes.int32),
+        (tensor_shape.TensorShape([None]), tensor_shape.TensorShape([None])))
+    padded_dataset = dataset.padded_batch(
+        2, padded_shapes=([None], [None]), padding_values=('', 0))
+    next_element = self.getNext(padded_dataset)
+    self.evaluate(next_element())
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPaddedBatchDatasetShapeSpecifications(self):
+    int_placeholder = array_ops.placeholder(dtypes.int32)
+    float_placeholder = array_ops.placeholder(dtypes.float32)
+    string_placeholder = array_ops.placeholder(dtypes.string)
+    input_dataset = dataset_ops.Dataset.from_tensors(
+        (int_placeholder, float_placeholder, string_placeholder))
+
+    # Test different ways of specifying the `padded_shapes` argument.
+    dynamic_padding_from_tensor_shapes = input_dataset.padded_batch(
+        32,
+        padded_shapes=(tensor_shape.TensorShape([None]),
+                       tensor_shape.TensorShape([None, None]),
+                       tensor_shape.TensorShape([37])))
+    dynamic_padding_from_lists = input_dataset.padded_batch(
+        32, padded_shapes=([None], [None, None], [37]))
+    dynamic_padding_from_lists_with_minus_one = input_dataset.padded_batch(
+        32, padded_shapes=([-1], [-1, -1], [37]))
+    dynamic_padding_from_tensors = input_dataset.padded_batch(
+        32,
+        padded_shapes=(constant_op.constant([-1], dtype=dtypes.int64),
+                       constant_op.constant([-1, -1], dtype=dtypes.int64),
+                       constant_op.constant([37], dtype=dtypes.int64)))
+
+    for dataset in [
+        dynamic_padding_from_tensor_shapes, dynamic_padding_from_lists,
+        dynamic_padding_from_lists_with_minus_one, dynamic_padding_from_tensors
+    ]:
+      dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset)
+      self.assertEqual([None, None], dataset_output_shapes[0].as_list())
+      self.assertEqual([None, None, None], dataset_output_shapes[1].as_list())
+      self.assertEqual([None, 37], dataset_output_shapes[2].as_list())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchSparseError(self):
+
+    st = sparse_tensor.SparseTensorValue(
+        indices=[[0, 0]], values=([42]), dense_shape=[1, 1])
+
+    with self.assertRaisesRegex(
+        TypeError, r'`padded_batch` is only supported for '
+        r'datasets that produce tensor elements but type spec of elements in '
+        r'the input dataset is not a subclass of TensorSpec: '
+        r'`SparseTensorSpec.*`\.$'):
+      _ = dataset_ops.Dataset.from_tensors(st).repeat(10).padded_batch(10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchRaggedError(self):
+
+    rt = ragged_tensor_value.RaggedTensorValue(
+        np.array([0, 42]), np.array([0, 2], dtype=np.int64))
+
+    with self.assertRaisesRegex(
+        TypeError, r'`padded_batch` is only supported for '
+        r'datasets that produce tensor elements but type spec of elements in '
+        r'the input dataset is not a subclass of TensorSpec: '
+        r'`RaggedTensorSpec.*`\.$'):
+      _ = dataset_ops.Dataset.from_tensors(rt).repeat(10).padded_batch(10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchDatasetsError(self):
+
+    ds = dataset_ops.Dataset.range(10).map(
+        lambda x: dataset_ops.Dataset.range(1))
+
+    with self.assertRaisesRegex(
+        TypeError, r'`padded_batch` is not supported for datasets of datasets'):
+      _ = ds.padded_batch(3)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorWrongRank(self):
+    with self.assertRaisesRegex(
+        ValueError, r'The padded shape \(1,\) is not compatible with the '
+        r'shape \(\) of the corresponding input component.'):
+      _ = dataset_ops.Dataset.range(10).padded_batch(5, padded_shapes=[1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorTooSmall(self):
+    with self.assertRaisesRegex(
+        ValueError, r'The padded shape \(1,\) is not compatible with the '
+        r'shape \(3,\) of the corresponding input component.'):
+      _ = dataset_ops.Dataset.from_tensors([1, 2, 3]).padded_batch(
+          5, padded_shapes=[1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorShapeNotRank1(self):
+    with self.assertRaisesRegex(
+        ValueError, r'Padded shape .* must be a `tf.int64` vector tensor, '
+        r'but its shape was \(2, 2\).'):
+      _ = dataset_ops.Dataset.from_tensors([1, 2, 3]).padded_batch(
+          5, padded_shapes=[[1, 1], [1, 1]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorShapeNotInt(self):
+    with self.assertRaisesRegex(
+        TypeError, r'Padded shape .* must be a `tf.int64` vector tensor, '
+        r'but its element type was float32.'):
+      _ = dataset_ops.Dataset.from_tensors([1, 2, 3]).padded_batch(
+          5, padded_shapes=constant_op.constant([1.5, 2., 3.]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorWrongRankFromTensor(self):
+    with self.assertRaisesRegex(
+        ValueError, r'The padded shape \(1,\) is not compatible with the '
+        r'shape \(\) of the corresponding input component.'):
+      shape_as_tensor = constant_op.constant([1], dtype=dtypes.int64)
+      _ = dataset_ops.Dataset.range(10).padded_batch(
+          5, padded_shapes=shape_as_tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchShapeErrorDefaultShapeWithUnknownRank(self):
+    with self.assertRaisesRegex(ValueError, r'`padded_shapes`.*unknown rank'):
+      ds = dataset_ops.Dataset.from_generator(
+          lambda: iter([1, 2, 3]), output_types=dtypes.int32)
+      ds.padded_batch(2)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPaddedBatchShapeErrorPlaceholder(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        r'The padded shape \((\?|None), (\?|None)\) is not compatible with the '
+        r'shape \(\) of the corresponding input component.'):
+      shape_as_tensor = array_ops.placeholder(dtypes.int64, shape=[2])
+      _ = dataset_ops.Dataset.range(10).padded_batch(
+          5, padded_shapes=shape_as_tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPaddedBatchBfloat16(self):
+    ds = dataset_ops.Dataset.range(5)
+    ds = ds.map(lambda x: math_ops.cast(x, dtypes.bfloat16))
+    ds = ds.padded_batch(10)
+    self.assertDatasetProduces(
+        ds, expected_output=[[0.0, 1.0, 2.0, 3.0, 4.0]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDefaultPaddedValueShapes(self):
+
+    def fill(x):
+      return array_ops.fill([x], x)
+
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4]).map(fill),
+         dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4]).map(fill)))
+    dataset = dataset.padded_batch(batch_size=2, padding_values=-1)
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[([[1, -1], [2, 2]], [[1, -1], [2, 2]]),
+                         ([[3, 3, 3, -1], [4, 4, 4, 4]], [[3, 3, 3, -1],
+                                                          [4, 4, 4, 4]])])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.range(5).padded_batch(5, name='padded_batch')
+    self.assertDatasetProduces(dataset, [list(range(5))])
+
+
+class PaddedBatchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+
+    def build_dataset(seq_lens):
+      dataset = dataset_ops.Dataset.from_tensor_slices(seq_lens).map(
+          lambda x: array_ops.fill([x], x)).padded_batch(
+              batch_size=4, padded_shapes=[-1])
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      dataset = dataset.with_options(options)
+      return dataset
+
+    seq_lens = np.random.randint(1, 20, size=(32,)).astype(np.int32)
+    verify_fn(self, lambda: build_dataset(seq_lens), num_outputs=8)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testNonDefaultPadding(self, verify_fn):
+
+    def build_dataset(seq_lens):
+
+      def fill_tuple(x):
+        filled = array_ops.fill([x], x)
+        return (filled, string_ops.as_string(filled))
+
+      padded_shape = [-1]
+      return dataset_ops.Dataset.from_tensor_slices(seq_lens).map(
+          fill_tuple).padded_batch(
+              batch_size=4,
+              padded_shapes=(padded_shape, padded_shape),
+              padding_values=(-1, '<end>'))
+
+    seq_lens = np.random.randint(1, 20, size=(32,)).astype(np.int32)
+    verify_fn(self, lambda: build_dataset(seq_lens), num_outputs=8)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/data/kernel_tests/placement_test.py

@@ -0,0 +1,294 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tf.data placement within tf.functions."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class PlacementTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for tf.data placement within tf.functions.
+
+  Specifically, tf.data dataset tensors cannot be copied between devices. These
+  tests verify the ops are placed in a way that avoids this.
+  """
+
+  def setUp(self):
+    super(PlacementTest, self).setUp()
+    # Grappler optimizations can affect whether the placement issues occur,
+    # since they may inadvertently rewrite nodes and edges in a way that removes
+    # cross-device copies.
+    config.set_optimizer_experimental_options({"disable_meta_optimizer": True})
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testWhileWithCapturedDataset(self):
+    dataset = dataset_ops.Dataset.range(10)
+
+    @def_function.function
+    def f():
+      total = constant_op.constant(0, dtypes.int64)
+      for _ in math_ops.range(1):
+        for elem in dataset:
+          total += elem
+      return total
+
+    self.assertEqual(f().numpy(), 45)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testWhile(self):
+
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(10)
+      total = constant_op.constant(0, dtypes.int64)
+      for _ in math_ops.range(1):
+        for elem in dataset:
+          total += elem
+      return total
+
+    self.assertEqual(f().numpy(), 45)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCondWithPlacement(self):
+    # When the cond op is explicitly placed, there shouldn't be cross-device
+    # copies.
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(10)
+
+      def fn():
+        return dataset.map(lambda x: x+1)
+
+      c = constant_op.constant(2)
+      with ops.device("/cpu:0"):
+        a = cond.cond(math_ops.equal(c, 2), fn, fn)
+        iterator = iter(a)
+        nxt = next(iterator)
+      return nxt
+
+    self.assertEqual(f().numpy(), 1)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCondWithColocation(self):
+    # When the cond op is colocated with the dataset, there shouldn't be
+    # cross-device copies.
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(8)
+
+      def fn():
+        return dataset.map(lambda x: x+1)
+
+      c = constant_op.constant(2)
+      with ops.colocate_with(dataset._variant_tensor):  # pylint:disable=protected-access
+        a = cond.cond(math_ops.equal(c, 2), fn, fn)
+        iterator = iter(a)
+        nxt = next(iterator)
+      return nxt
+
+    self.assertEqual(f().numpy(), 1)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCond(self):
+
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(8)
+      c = constant_op.constant(2)
+      a = cond.cond(
+          math_ops.equal(c, 2),
+          lambda: dataset.map(lambda x: x + 1),
+          lambda: dataset.map(lambda x: x + 2),
+      )
+      return next(iter(a))
+
+    self.assertEqual(f().numpy(), 1)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testId(self):
+    # Ideally, placer should know that Identity(dataset) should be on the same
+    # device as the dataset.
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(10)
+      dataset = array_ops.identity(dataset)
+      return dataset
+    f()
+
+  @combinations.generate(test_base.default_test_combinations())
+  @test_util.run_gpu_only
+  def testFunctionCall(self):
+    # Ideally, placer should know that Call(dataset) should be on the same
+    # device as the dataset. Create a function that could be place don the GPU,
+    # but a Dataset that cannot.
+    @def_function.function
+    def test_call(dataset):
+      return dataset.reduce(0, lambda s, _: s + 1)
+
+    @def_function.function
+    def f():
+      dataset = dataset_ops.Dataset.range(10)
+      return test_call(dataset)
+
+    self.assertEqual(self.evaluate(f()), 10)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  @test_util.run_gpu_only
+  def testIteratorOnDeviceEagerMode(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
+    iterator = iter(dataset)
+    data = next(iterator)
+    optional_data = iterator.get_next_as_optional()
+
+    self.assertIn("gpu:0", dataset._variant_tensor.device.lower())
+    self.assertIn("gpu:0", iterator._iterator_resource.device.lower())
+    self.assertIn("gpu:0", data.device.lower())
+    self.assertIn("gpu:0", optional_data.get_value().device.lower())
+    self.assertIn("gpu:0", optional_data.has_value().device.lower())
+
+  # There are HostMemory constraints on AnonymousIteratorV2 and
+  # DeleteIterator kernels on TPU but not on GPU. This is intentional because
+  # when running AnonymousIteratorV2 in a function
+  #
+  # - If the op is placed on GPU, the variant _Retval is placed on GPU.
+  # - However, if the op is placed on TPU, the variant _Retval is placed on
+  #   CPU.
+  #
+  # So if were to add HostMemory constraints to the GPU kernels it would lead
+  # to variant device copy errors.
+  #
+  # TODO(b/204231062): Unify behavior across GPU and TPU.
+  @combinations.generate(test_base.eager_only_combinations())
+  @test_util.run_gpu_only
+  def testCreateIteratorInFuncOnGpu(self):
+
+    @def_function.function
+    def create_iter():
+      return gen_dataset_ops.anonymous_iterator_v2(
+          output_types=[dtypes.float32], output_shapes=[[]])
+
+    create_iter()
+
+  @combinations.generate(test_base.graph_only_combinations())
+  @test_util.run_gpu_only
+  def testIteratorOnDeviceGraphModeOneShotIterator(self):
+    self.skipTest("TODO(b/169429285): tf.data.Dataset.make_one_shot_iterator "
+                  "does not support GPU placement.")
+
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    data = iterator.get_next()
+    optional_data = iterator.get_next_as_optional()
+
+    with ops.colocate_with(dataset._variant_tensor):
+      dataset_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(dataset_device))
+
+    with ops.colocate_with(iterator._iterator_resource):
+      iterator_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(iterator_device))
+
+    with ops.colocate_with(data):
+      data_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(data_device))
+
+    with ops.colocate_with(optional_data.get_value()):
+      get_value_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(get_value_device))
+
+    with ops.colocate_with(optional_data.has_value()):
+      has_value_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(has_value_device))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  @test_util.run_gpu_only
+  def testIteratorOnDeviceGraphModeInitializableIterator(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    data = iterator.get_next()
+    optional_data = iterator.get_next_as_optional()
+
+    with ops.colocate_with(dataset._variant_tensor):
+      dataset_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(dataset_device))
+
+    with ops.colocate_with(iterator._iterator_resource):
+      iterator_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(iterator_device))
+
+    with ops.colocate_with(data):
+      data_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(data_device))
+
+    with ops.colocate_with(optional_data.get_value()):
+      get_value_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(get_value_device))
+
+    with ops.colocate_with(optional_data.has_value()):
+      has_value_device = test_ops.device_placement_op()
+    self.assertIn(b"GPU:0", self.evaluate(has_value_device))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  @test_util.run_gpu_only
+  def testIterDatasetEagerModeWithExplicitDevice(self):
+
+    @def_function.function
+    def comp():
+      value = constant_op.constant(0, dtype=dtypes.int64)
+      for d in iter(dataset_ops.Dataset.range(10)):
+        value += d
+      return value
+
+    with ops.device("/gpu:0"):
+      result = comp()
+    self.assertEqual(result.numpy(), 45)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  @test_util.run_gpu_only
+  def testFunctionInliningColocation(self):
+
+    @def_function.function
+    def f(ds):
+      return next(iter(ds))
+
+    @def_function.function
+    def g():
+      dataset = dataset_ops.Dataset.range(10)
+      return f(dataset)
+
+    with ops.device("/gpu:0"):
+      self.assertEqual(self.evaluate(g()), 0)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/prefetch_test.py

@@ -0,0 +1,170 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.prefetch()`."""
+import time
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import prefetch_op
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import script_ops
+from tensorflow.python.platform import test
+
+
+class PrefetchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(buffer_size=[-1, None, 0, 42])))
+  def testBufferSize(self, buffer_size):
+    dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(buffer_size=[0, 1, 2, 42])))
+  def testPrefetching(self, buffer_size):
+    dataset = dataset_ops.Dataset.range(1000)
+
+    calls = 0
+
+    @script_ops.eager_py_func(Tout=[dtypes.int64])
+    def map_fn(x):
+      nonlocal calls
+      calls += 1
+      return x
+
+    dataset = dataset.map(map_fn)
+    dataset = dataset.prefetch(buffer_size=buffer_size)
+    it = iter(dataset)
+    for _ in range(10):
+      next(it)
+
+    # Wait for the prefetch buffer to fill up.
+    while calls != 10+buffer_size:
+      time.sleep(0.1)
+    # Wait some extra time to make sure the prefetch buffer isn't fetching more
+    # elements than it should.
+    time.sleep(0.5)
+    self.assertEqual(calls, 10+buffer_size)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(buffer_size=[-2, -42])))
+  def testInvalidBufferSize(self, buffer_size):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size)
+      self.evaluate(dataset._variant_tensor)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              buffer_size=[-1, None, 0, 42], slack_period=[1, 8])))
+  def testPrefetchWithSlack(self, buffer_size, slack_period):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = prefetch_op._PrefetchDataset(  # pylint: disable=protected-access
+        dataset, buffer_size, slack_period=slack_period)
+    self.assertDatasetProduces(dataset, expected_output=range(100))
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testPrefetchCancellation(self):
+
+    def map_py_fn(x):
+      while x > -1:
+        x = x * 1
+      return x
+
+    dataset = dataset_ops.Dataset.range(10).map(map_py_fn).prefetch(3)
+    get_next = self.getNext(dataset)
+
+    with self.cached_session() as sess:
+      thread = self.checkedThread(self.assert_op_cancelled, args=(get_next(),))
+      thread.start()
+      time.sleep(2)
+      sess.close()
+      thread.join()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).prefetch(1, name="prefetch")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class PrefetchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                             parameterized.TestCase):
+
+  def build_dataset(self, options=None):
+    dataset = dataset_ops.Dataset.range(100).prefetch(10)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self.build_dataset(options), num_outputs=100)
+
+
+class PrefetchRandomAccessTest(test_base.DatasetTestBase,
+                               parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 10, 11])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=5)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-2, 0, 1])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).prefetch(buffer_size=5)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(elements=[10, 50, 100], buffer_size=[0, 5, 10])))
+  def testMultipleCombinations(self, elements, buffer_size):
+    dataset = dataset_ops.Dataset.range(elements).prefetch(
+        buffer_size=buffer_size)
+    len_dataset = self.evaluate(dataset.cardinality())
+    expected_output = np.arange(elements)
+    for i in range(len_dataset):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)), expected_output[i])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=len_dataset))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/ragged_batch_test.py

@@ -0,0 +1,171 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.ragged_batch`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+
+
+def _make_scalar_ds(nrows):
+  """Create a test dataset with scalar elements."""
+  return dataset_ops.Dataset.from_tensor_slices(np.arange(nrows))
+
+
+def _make_vector_ds(nrows):
+  """Create a test dataset with vector elements (of varying size)."""
+  return _make_scalar_ds(nrows).map(lambda x: array_ops.fill([x], x))
+
+
+def _make_matrix_ds1(nrows):
+  """Create a test dataset with matrix elements (of varying size)."""
+  return _make_scalar_ds(nrows).map(lambda x: array_ops.fill([x, 2], x))
+
+
+def _make_matrix_ds2(nrows):
+  """Create a test dataset with matrix elements (of varying size)."""
+  return _make_scalar_ds(nrows).map(lambda x: array_ops.fill([2, x], x))
+
+
+def _make_matrix_ds_fully_defined(nrows):
+  """Create a test dataset with matrix elements (of varying size)."""
+  return _make_scalar_ds(nrows).map(lambda x: array_ops.fill([2, 3], x))
+
+
+def _make_5dtensor_ds(nrows):
+  """Create a test dataset with matrix elements (of varying size)."""
+  return _make_scalar_ds(nrows).map(
+      lambda x: array_ops.fill([2, x, 3, 2*x, 4], x))
+
+
+def _make_ragged_ds(nrows):
+  """Create a test dataset with RaggedTensor elements (of varying size)."""
+  values = [[[i] * (i % 3) for i in range(j)] * (j % 3) for j in range(nrows)]
+  rt = ragged_factory_ops.constant(values)
+  return dataset_ops.Dataset.from_tensor_slices(rt)
+
+
+def _make_dict_ds(nrows):
+  """Create a test set with various element shapes."""
+  def transform(x):
+    return {
+        'shape=[]': ops.convert_to_tensor(x),
+        'shape=[x]': math_ops.range(x),
+        'shape=[x, 2]': array_ops.fill([x, 2], x),
+        'shape=[2, x]': array_ops.fill([2, x], x),
+        'shape=[2, x, 3, 2x, 4]': array_ops.fill([2, x, 3, 2*x, 4], x)
+    }
+  return _make_scalar_ds(nrows).map(transform)
+
+
+def _make_tuple_ds(nrows):
+  """Create a test set with various element shapes."""
+  def transform(x):
+    return (ops.convert_to_tensor(x),
+            math_ops.range(x),
+            array_ops.fill([x, 2], x))
+  return _make_scalar_ds(nrows).map(transform)
+
+
+def _to_list(v):
+  return v.to_list() if hasattr(v, 'to_list') else v.tolist()
+
+
+class RaggedBatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              make_dataset=[
+                  _make_scalar_ds, _make_vector_ds, _make_matrix_ds1,
+                  _make_matrix_ds2, _make_ragged_ds, _make_5dtensor_ds,
+                  _make_dict_ds, _make_tuple_ds, _make_matrix_ds_fully_defined,
+              ],
+              nrows=[0, 20, 23],
+              batch_size=[4],
+              drop_remainder=[True, False])))
+  def testBasic(self, make_dataset, nrows, batch_size, drop_remainder):
+    dataset = make_dataset(nrows)
+
+    # Get the unbatched rows (so we can check expected values).
+    get_next = self.getNext(dataset)
+    rows = [nest.map_structure(_to_list, self.evaluate(get_next()))
+            for _ in range(nrows)]
+
+    # Batch the dataset, and check that batches match slices from `rows`.
+    batched_dataset = dataset.ragged_batch(batch_size, drop_remainder)
+    get_next = self.getNext(batched_dataset)
+    for start_row in range(0, nrows, batch_size):
+      end_row = start_row + batch_size
+      if end_row > nrows and drop_remainder:
+        break
+      end_row = min(end_row, nrows)
+      result = self.evaluate(get_next())
+
+      # Use nest for potentially nested datasets.
+      nest.map_structure_up_to(
+          result, lambda a, *b: self.assertAllEqual(a, list(b)),
+          result, *rows[start_row:end_row])
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithStructuredElements(self):
+    nrows = 20
+    batch_size = 4
+
+    def make_structure(x):
+      return {
+          'dense':
+              array_ops.fill([x], x),
+          'ragged':
+              ragged_concat_ops.stack(
+                  [array_ops_stack.stack([x]),
+                   array_ops_stack.stack([x, x])]),
+          'sparse':
+              sparse_tensor.SparseTensor([[x]], [x], [100])
+      }
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(np.arange(nrows))
+    dataset = dataset.map(make_structure)
+    dataset = dataset.ragged_batch(batch_size)
+    get_next = self.getNext(dataset)
+
+    for i in range(0, nrows, batch_size):
+      result = self.evaluate(get_next())
+      rows = range(i, i + batch_size)
+      self.assertAllEqual(result['dense'], [[r] * r for r in rows])
+      self.assertAllEqual(result['ragged'], [[[r], [r, r]] for r in rows])
+      self.assertAllEqual(result['sparse'].indices, list(enumerate(rows)))
+      self.assertAllEqual(result['sparse'].values, rows)
+      self.assertAllEqual(result['sparse'].dense_shape, [4, 100])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/data/kernel_tests/random_test.py

@@ -0,0 +1,179 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.random()`."""
+import warnings
+
+from absl.testing import parameterized
+
+from tensorflow.python import tf2
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import random_seed
+from tensorflow.python.platform import test
+
+
+class RandomTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(global_seed=[None, 10], local_seed=[None, 20])))
+  def testDeterminism(self, global_seed, local_seed):
+    expect_determinism = (global_seed is not None) or (local_seed is not None)
+
+    random_seed.set_random_seed(global_seed)
+    ds = dataset_ops.Dataset.random(seed=local_seed).take(10)
+
+    output_1 = self.getDatasetOutput(ds, requires_initialization=True)
+    ds = self.graphRoundTrip(ds)
+    output_2 = self.getDatasetOutput(ds, requires_initialization=True)
+
+    if expect_determinism:
+      self.assertEqual(output_1, output_2)
+    else:
+      # Technically not guaranteed since the two randomly-chosen int64 seeds
+      # could match, but that is sufficiently unlikely (1/2^128 with perfect
+      # random number generation).
+      self.assertNotEqual(output_1, output_2)
+
+  @combinations.generate(
+      combinations.times(test_base.graph_only_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testRerandomizeEachIterationEpochsIgnored(self, rerandomize):
+    with warnings.catch_warnings(record=True) as w:
+      dataset = dataset_ops.Dataset.random(
+          seed=42,
+          rerandomize_each_iteration=rerandomize,
+          name="random").take(10)
+    first_epoch = self.getDatasetOutput(dataset, requires_initialization=True)
+    second_epoch = self.getDatasetOutput(dataset, requires_initialization=True)
+    if rerandomize:
+      if not tf2.enabled() and rerandomize:
+        found_warning = False
+        for warning in w:
+          if ("In TF 1, the `rerandomize_each_iteration=True` option" in
+              str(warning)):
+            found_warning = True
+            break
+        self.assertTrue(found_warning)
+
+    self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testRerandomizeEachIterationEpochs(self, rerandomize):
+    dataset = dataset_ops.Dataset.random(
+        seed=42, rerandomize_each_iteration=rerandomize, name="random").take(10)
+    first_epoch = self.getDatasetOutput(dataset)
+    second_epoch = self.getDatasetOutput(dataset)
+
+    if rerandomize:
+      self.assertEqual(first_epoch == second_epoch,
+                       not rerandomize or rerandomize is None)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testRerandomizeRepeatEpochs(self, rerandomize):
+    dataset = dataset_ops.Dataset.random(
+        seed=42, rerandomize_each_iteration=rerandomize, name="random").take(10)
+    dataset = dataset.repeat(2)
+    next_element = self.getNext(dataset, requires_initialization=True)
+    first_epoch = []
+    for _ in range(10):
+      first_epoch.append(self.evaluate(next_element()))
+    second_epoch = []
+    for _ in range(10):
+      second_epoch.append(self.evaluate(next_element()))
+
+    if rerandomize:
+      self.assertEqual(first_epoch == second_epoch,
+                       not rerandomize or rerandomize is None)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(
+      combinations.times(test_base.v2_eager_only_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testRerandomizeInsideFunction(self, rerandomize):
+
+    @def_function.function
+    def make_dataset():
+      dataset = dataset_ops.Dataset.random(
+          seed=42,
+          rerandomize_each_iteration=rerandomize,
+          name="random").take(10)
+      return dataset
+
+    dataset = make_dataset()
+    first_epoch = self.getDatasetOutput(dataset)
+    second_epoch = self.getDatasetOutput(dataset)
+
+    if rerandomize:
+      self.assertEqual(first_epoch == second_epoch,
+                       not rerandomize or rerandomize is None)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.random(
+        seed=42, name="random").take(1).map(lambda _: 42)
+    self.assertDatasetProduces(dataset, expected_output=[42],
+                               requires_initialization=True)
+
+
+class RandomCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                           parameterized.TestCase):
+
+  def _build_random_dataset(
+      self,
+      num_elements=10,
+      seed=None,
+      rerandomize_each_iteration=None):
+    dataset = dataset_ops.Dataset.random(
+        seed=seed, rerandomize_each_iteration=rerandomize_each_iteration)
+    # Checkpoint tests need the test dataset to be finite whereas `random` is
+    # infinite. Use `take` to limit the number of elements.
+    return dataset.take(num_elements)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              rerandomize_each_iteration=[True, False])))
+
+  def test(self, verify_fn, rerandomize_each_iteration):
+    seed = 55
+    num_elements = 10
+    # pylint: disable=g-long-lambda
+    verify_fn(
+        self,
+        lambda: self._build_random_dataset(
+            seed=seed,
+            num_elements=num_elements,
+            rerandomize_each_iteration=rerandomize_each_iteration),
+        num_elements)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/range_test.py

@@ -0,0 +1,222 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.range()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class RangeTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStop(self, output_type):
+    stop = 5
+    dataset = dataset_ops.Dataset.range(stop, output_type=output_type)
+    expected_output = np.arange(stop, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStartStop(self, output_type):
+    start, stop = 2, 5
+    dataset = dataset_ops.Dataset.range(start, stop, output_type=output_type)
+    expected_output = np.arange(start, stop, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStartStopStep(self, output_type):
+    start, stop, step = 2, 10, 2
+    dataset = dataset_ops.Dataset.range(
+        start, stop, step, output_type=output_type)
+    expected_output = np.arange(
+        start, stop, step, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testZeroStep(self, output_type):
+    start, stop, step = 2, 10, 0
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(
+          start, stop, step, output_type=output_type)
+      self.evaluate(dataset._variant_tensor)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testNegativeStep(self, output_type):
+    start, stop, step = 2, 10, -1
+    dataset = dataset_ops.Dataset.range(
+        start, stop, step, output_type=output_type)
+    expected_output = np.arange(
+        start, stop, step, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStopLessThanStart(self, output_type):
+    start, stop = 10, 2
+    dataset = dataset_ops.Dataset.range(start, stop, output_type=output_type)
+    expected_output = np.arange(start, stop, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStopLessThanStartWithPositiveStep(self, output_type):
+    start, stop, step = 10, 2, 2
+    dataset = dataset_ops.Dataset.range(
+        start, stop, step, output_type=output_type)
+    expected_output = np.arange(
+        start, stop, step, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(output_type=[
+              dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+          ])))
+  def testStopLessThanStartWithNegativeStep(self, output_type):
+    start, stop, step = 10, 2, -1
+    dataset = dataset_ops.Dataset.range(
+        start, stop, step, output_type=output_type)
+    expected_output = np.arange(
+        start, stop, step, dtype=output_type.as_numpy_dtype)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+    self.assertEqual(output_type, dataset_ops.get_legacy_output_types(dataset))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.range(5, name="range")
+    self.assertDatasetProduces(dataset, list(range(5)))
+
+
+class RangeCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                          parameterized.TestCase):
+
+  def _build_range_dataset(self, start, stop, options=None):
+    dataset = dataset_ops.Dataset.range(start, stop)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    start = 2
+    stop = 10
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self._build_range_dataset(start, stop, options),
+              stop - start)
+
+
+class RangeRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 2, 3])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(2)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 0])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.range(0)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testBasic(self):
+    dataset = dataset_ops.Dataset.range(10)
+    for i in range(10):
+      self.assertEqual(self.evaluate(random_access.at(dataset, index=i)), i)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              start=[-1, 0, 5],
+              stop=[-5, 0, 10],
+              step=[-3, 1, 5],
+              output_type=[
+                  dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64
+              ])))
+  def testMultipleCombinations(self, start, stop, step, output_type):
+    dataset = dataset_ops.Dataset.range(
+        start, stop, step, output_type=output_type)
+    expected_output = np.arange(
+        start, stop, step, dtype=output_type.as_numpy_dtype)
+    len_dataset = self.evaluate(dataset.cardinality())
+    for i in range(len_dataset):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)), expected_output[i])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=len_dataset))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/rebatch_test.py

@@ -0,0 +1,362 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.rebatch()`."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+def _flat_shapes(dataset):
+  return [
+      ts.as_list()
+      for ts in nest.flatten(dataset_ops.get_legacy_output_shapes(dataset))
+  ]
+
+
+class RebatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  ##############################################################################
+  # The following tests exercise our static computation of output_shapes.
+  ##############################################################################
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShapeInferenceNotAllBatchSizesEqual(self):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 1, 1])
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testShapeInferenceInputBatchDimDivisible(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(
+        batch_size=[2, 2], drop_remainder=drop_remainder)
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testShapeInferenceInputBatchDimUnknown(self):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=False)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=False)
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testShapeInferenceInputBatchDimUnknownWithDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=False)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=True)
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testShapeInferenceInputBatchDimIndivisible(self):
+    dataset = dataset_ops.Dataset.range(10).batch(5, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=False)
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testShapeInferenceInputBatchDimIndivisibleWithDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(10).batch(5, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=True)
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+  ##############################################################################
+  # The following tests check `tf.data.Dataset.rebatch`'s output.
+  ##############################################################################
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testBasic(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2],
+                                        drop_remainder=drop_remainder)
+
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+    expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testPartialBatch(self):
+    dataset = dataset_ops.Dataset.range(5).batch(4, drop_remainder=False)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=False)
+
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+    expected_output = [[0, 1], [2, 3], [4]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testPartialBatchWithDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(5).batch(4, drop_remainder=False)
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2], drop_remainder=True)
+
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+    expected_output = [[0, 1], [2, 3]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testBatchSizeGreaterThanOriginal(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(12).batch(4, drop_remainder=False)
+    rebatched_dataset = dataset.rebatch(batch_size=[6],
+                                        drop_remainder=drop_remainder)
+
+    expected_output = [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              batch_size=[2, 3, 4], drop_remainder=[True, False]
+          ),
+      )
+  )
+  def testBatchSizeEqualToOriginal(self, batch_size, drop_remainder):
+    # `drop_remainder` needs to be `False` in `rebatch` call
+    # so that the remainder batch is preserved.
+    #
+    # For example:
+    # d = range(3).batch(2, drop_remainder=True)
+    # d2 = d.rebatch(2, drop_remainder=True)
+    # d becomes [[0, 1], [2]] and d2 becomes [[0, 1]],
+    # which is a mismatch we do not want.
+
+    dataset = dataset_ops.Dataset.range(11).batch(
+        batch_size, drop_remainder=drop_remainder
+    )
+    expected_output = self.getDatasetOutput(dataset)
+    rebatched_dataset = dataset.rebatch(
+        batch_size=batch_size, drop_remainder=False
+    )
+
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(drop_remainder=[True, False]),
+      )
+  )
+  def testEmptySplits(self, drop_remainder):
+    # It's possible for splits to be empty if the batch size is smaller than
+    # the number of replicas. Here, we use an example with batch_size == 4
+    # and num_replicas == 5.
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[1, 1, 1, 1, 0],
+                                        drop_remainder=drop_remainder)
+
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+    expected_output = [[0], [1], [2], [3], [], [4], [5], [6], [7], []]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testEmptyFirstSplits(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[0, 1],
+                                        drop_remainder=drop_remainder)
+
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+    # We have an extra element at the end because if the desired batch size is
+    # zero, then we never read any inputs from the input_dataset at all, so we
+    # will keep producing empty outputs until we reach a non zero desired batch
+    # size split.
+    expected_output = [[], [0], [], [1], [], [2], [], [3], [], [4], [], [5], [],
+                       [6], [], [7], []]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testEmptyLastSplits(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=[1, 0],
+                                        drop_remainder=drop_remainder)
+
+    expected_shapes = [[None]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+    expected_output = [[0], [], [1], [], [2], [], [3], [], [4], [], [5], [],
+                       [6], [], [7], []]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testEmptyTensors(self, drop_remainder):
+    """Tests empty tensors case.
+
+    Args:
+      drop_remainder: whether to drop the remainder.
+
+    The implementation of rebatch might move the input data.
+    This test ensures the empty buffer is handled correctly.
+    """
+    new_batch_size = 4
+    dataset = dataset_ops.Dataset.range(8)
+    dataset = dataset.map(lambda x: array_ops.reshape((), (5, 0)))
+    dataset = dataset.batch(2)
+    rebatched_dataset = dataset.rebatch(
+        batch_size=new_batch_size, drop_remainder=drop_remainder
+    )
+
+    expected_output = [
+        array_ops.reshape((), (new_batch_size, 5, 0))
+        for _ in range(8 // new_batch_size)
+    ]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(drop_remainder=[True, False]),
+      )
+  )
+  def testScalarBatchSizeInput(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch(batch_size=2,
+                                        drop_remainder=drop_remainder)
+
+    expected_shapes = [[2]]
+    self.assertEqual(expected_shapes, _flat_shapes(rebatched_dataset))
+
+    expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleBatches(self):
+    dataset = dataset_ops.Dataset.range(16).batch(
+        2, drop_remainder=True).batch(
+            4, drop_remainder=True)
+    self.assertEqual([[4, 2]], _flat_shapes(dataset))
+
+    rebatched_dataset = dataset.rebatch([2, 2])
+    self.assertEqual([[2, 2]], _flat_shapes(rebatched_dataset))
+    # Each element is a list of 2 elements where each element is a list of 2.
+    expected_output = [[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8, 9], [10, 11]],
+                       [[12, 13], [14, 15]]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedDictionaryOutput(self):
+
+    def map_fn(x):
+      return {"a": x, "b": {"c": x + 1}}
+
+    dataset = dataset_ops.Dataset.range(8).map(map_fn).batch(
+        4, drop_remainder=True)
+    rebatched_dataset = dataset.rebatch([2, 2])
+    self.assertEqual([[2], [2]], _flat_shapes(rebatched_dataset))
+
+    expected_output = [{
+        "a": [0, 1],
+        "b": {
+            "c": [1, 2]
+        }
+    }, {
+        "a": [2, 3],
+        "b": {
+            "c": [3, 4]
+        }
+    }, {
+        "a": [4, 5],
+        "b": {
+            "c": [5, 6]
+        }
+    }, {
+        "a": [6, 7],
+        "b": {
+            "c": [7, 8]
+        }
+    }]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testRaggedDataset(self, drop_remainder):
+    # Set up a dataset that produces ragged tensors with a static batch size.
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        ragged_tensor.RaggedTensor.from_row_lengths(
+            list(range(10)), [1, 2, 3, 4]))
+    # The map changes the internal representation of the ragged tensor.
+    # This test will fail if we don't normalize the tensor representation.
+    dataset = dataset.batch(4, drop_remainder=True).map(lambda x: x)
+
+    rebatched_dataset = dataset.rebatch(batch_size=[2, 2])
+
+    expected_output = [
+        ragged_tensor.RaggedTensor.from_row_lengths(list(range(3)), [1, 2]),
+        ragged_tensor.RaggedTensor.from_row_lengths(list(range(3, 10)), [3, 4]),
+    ]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoneDataset(self):
+    # Some datasets, e.g. datasets with None tensors, have components without
+    # output shapes. Test that this doesn't break rebatching shape inference
+    # logic.
+    dataset = dataset_ops.Dataset.range(4)
+    dataset = dataset.map(lambda x: (x, None))
+    dataset = dataset.batch(4, drop_remainder=True)
+    _ = dataset.rebatch(batch_size=[2, 2])
+
+
+class RebatchDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                   parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+
+    def build_dataset(num_elements, batch_size):
+      dataset = dataset_ops.Dataset.range(num_elements)
+      dataset_batched = dataset.batch(2 * batch_size, drop_remainder=True)
+      return dataset_batched.rebatch(batch_size=[batch_size, batch_size])
+
+    verify_fn(self, lambda: build_dataset(64, 8), num_outputs=8)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/reduce_test.py

@@ -0,0 +1,271 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.reduce()`."""
+import time
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class ReduceTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSum(self):
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1)
+      result = ds.reduce(np.int64(0), lambda x, y: x + y)
+      self.assertEqual(((i + 1) * i) // 2, self.evaluate(result))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSumTuple(self):
+
+    def reduce_fn(state, value):
+      v1, v2 = value
+      return state + v1 + v2
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1)
+      ds = dataset_ops.Dataset.zip((ds, ds))
+      result = ds.reduce(constant_op.constant(0, dtype=dtypes.int64), reduce_fn)
+      self.assertEqual(((i + 1) * i), self.evaluate(result))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSumAndCount(self):
+
+    def reduce_fn(state, value):
+      s, c = state
+      return s + value, c + 1
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1)
+      result = ds.reduce((constant_op.constant(0, dtype=dtypes.int64),
+                          constant_op.constant(0, dtype=dtypes.int64)),
+                         reduce_fn)
+      s, c = self.evaluate(result)
+      self.assertEqual(((i + 1) * i) // 2, s)
+      self.assertEqual(i, c)
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testSquareUsingPlaceholder(self):
+    delta = array_ops.placeholder(dtype=dtypes.int64)
+
+    def reduce_fn(state, _):
+      return state + delta
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1)
+      result = ds.reduce(np.int64(0), reduce_fn)
+      with self.cached_session() as sess:
+        square = sess.run(result, feed_dict={delta: i})
+        self.assertEqual(i * i, square)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparse(self):
+
+    def reduce_fn(_, value):
+      return value
+
+    def make_sparse_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1])),
+          dense_shape=np.array([1, 1]))
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.from_tensors(make_sparse_fn(i+1))
+      result = ds.reduce(make_sparse_fn(0), reduce_fn)
+      self.assertValuesEqual(make_sparse_fn(i + 1), self.evaluate(result))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNested(self):
+
+    def reduce_fn(state, value):
+      state["dense"] += value["dense"]
+      state["sparse"] = value["sparse"]
+      return state
+
+    def make_sparse_fn(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1])),
+          dense_shape=np.array([1, 1]))
+
+    def map_fn(i):
+      return {"dense": math_ops.cast(i, dtype=dtypes.int64),
+              "sparse": make_sparse_fn(math_ops.cast(i, dtype=dtypes.int64))}
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1).map(map_fn)
+      result = ds.reduce(map_fn(0), reduce_fn)
+      result = self.evaluate(result)
+      self.assertEqual(((i + 1) * i) // 2, result["dense"])
+      self.assertValuesEqual(make_sparse_fn(i), result["sparse"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDatasetSideEffect(self):
+    counter_var = variables.Variable(0)
+
+    def increment_fn(x):
+      counter_var.assign_add(1)
+      return x
+
+    def dataset_fn():
+      return dataset_ops.Dataset.range(10).map(increment_fn)
+
+    def reduce_fn(state, value):
+      return state + value
+
+    @def_function.function
+    def fn():
+      _ = dataset_fn().reduce(np.int64(0), reduce_fn)
+      return "hello"
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    self.assertEqual(self.evaluate(counter_var), 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSideEffect(self):
+    counter_var = variables.Variable(0)
+
+    def dataset_fn():
+      return dataset_ops.Dataset.range(10)
+
+    def reduce_fn(state, value):
+      counter_var.assign_add(1)
+      return state + value
+
+    @def_function.function
+    def fn():
+      _ = dataset_fn().reduce(np.int64(0), reduce_fn)
+      return "hello"
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    self.assertEqual(self.evaluate(counter_var), 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAutomaticControlDependencies(self):
+    counter_var = variables.Variable(1)
+
+    def dataset_fn():
+      return dataset_ops.Dataset.range(1)
+
+    def reduce1_fn(state, value):
+      counter_var.assign(counter_var + 1)
+      return state + value
+
+    def reduce2_fn(state, value):
+      counter_var.assign(counter_var * 2)
+      return state + value
+
+    @def_function.function
+    def fn():
+      _ = dataset_fn().reduce(np.int64(0), reduce1_fn)
+      _ = dataset_fn().reduce(np.int64(0), reduce2_fn)
+      return "hello"
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    self.assertEqual(self.evaluate(counter_var), 4)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedAutomaticControlDependencies(self):
+    counter_var = variables.Variable(0)
+
+    def map_fn(x):
+      counter_var.assign_add(1)
+      return x
+
+    def dataset_fn():
+      return dataset_ops.Dataset.range(10).map(map_fn)
+
+    @def_function.function
+    def fn():
+      for _ in dataset_fn():
+        pass
+      return counter_var
+
+    self.evaluate(counter_var.initializer)
+    self.assertEqual(self.evaluate(fn()), 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStateOnGPU(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPUs available.")
+
+    state = constant_op.constant(0, dtype=dtypes.int64)
+
+    def reduce_fn(state, value):
+      with ops.device("/gpu:0"):
+        return state + value
+
+    for i in range(10):
+      ds = dataset_ops.Dataset.range(1, i + 1)
+      result = ds.reduce(state, reduce_fn)
+      self.assertEqual(((i + 1) * i) // 2, self.evaluate(result))
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testCancellation(self):
+    ds = dataset_ops.Dataset.from_tensors(1).repeat()
+    result = ds.reduce(0, lambda x, y: x + y)
+    with self.cached_session() as sess:
+      # The `result` op is guaranteed to not complete before cancelled because
+      # the dataset that is being reduced is infinite.
+      thread = self.checkedThread(self.assert_op_cancelled, args=(result,))
+      thread.start()
+      time.sleep(0.2)
+      sess.close()
+      thread.join()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidFunction(self):
+    ds = dataset_ops.Dataset.range(5)
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(ds.reduce(0, lambda _, __: ()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptions(self):
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.apply(testing.assert_next(["MapAndBatch"]))
+    dataset = dataset.map(lambda x: x * 2).batch(5)
+    self.evaluate(dataset.reduce(0, lambda state, value: state))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42)
+    self.assertEqual(
+        self.evaluate(
+            dataset.reduce(0, lambda state, value: value, name="reduce")), 42)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/rejection_resample_test.py

@@ -0,0 +1,172 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.rejection_resample()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class RejectionResampleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(initial_known=[True, False])))
+  def testDistribution(self, initial_known):
+    classes = np.random.randint(5, size=(10000,))  # Uniformly sampled
+    target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
+    initial_dist = [0.2] * 5 if initial_known else None
+    classes = math_ops.cast(classes, dtypes.int64)  # needed for Windows build.
+    dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
+        200, seed=21).map(lambda c: (c, string_ops.as_string(c))).repeat()
+
+    get_next = self.getNext(
+        dataset.rejection_resample(
+            target_dist=target_dist,
+            initial_dist=initial_dist,
+            class_func=lambda c, _: c,
+            seed=27), requires_initialization=True)
+
+    returned = []
+    while len(returned) < 2000:
+      returned.append(self.evaluate(get_next()))
+
+    returned_classes, returned_classes_and_data = zip(*returned)
+    _, returned_data = zip(*returned_classes_and_data)
+    self.assertAllEqual([compat.as_bytes(str(c)) for c in returned_classes],
+                        returned_data)
+    total_returned = len(returned_classes)
+    class_counts = np.array(
+        [len([True for v in returned_classes if v == c]) for c in range(5)])
+    returned_dist = class_counts / total_returned
+    self.assertAllClose(target_dist, returned_dist, atol=1e-2)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(only_initial_dist=[True, False])))
+  def testEdgeCasesSampleFromInitialDataset(self, only_initial_dist):
+    init_dist = [0.5, 0.5]
+    target_dist = [0.5, 0.5] if only_initial_dist else [0.0, 1.0]
+    num_classes = len(init_dist)
+    # We don't need many samples to test that this works.
+    num_samples = 100
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(data_np)
+
+    # Reshape distribution.
+    dataset = dataset.rejection_resample(
+        class_func=lambda x: x, target_dist=target_dist, initial_dist=init_dist)
+
+    get_next = self.getNext(dataset)
+
+    returned = []
+    with self.assertRaises(errors.OutOfRangeError):
+      while True:
+        returned.append(self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRandomClasses(self):
+    init_dist = [0.25, 0.25, 0.25, 0.25]
+    target_dist = [0.0, 0.0, 0.0, 1.0]
+    num_classes = len(init_dist)
+    # We don't need many samples to test a dirac-delta target distribution.
+    num_samples = 100
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(data_np)
+
+    # Apply a random mapping that preserves the data distribution.
+    def _remap_fn(_):
+      return math_ops.cast(
+          random_ops.random_uniform([1]) * num_classes, dtypes.int32)[0]
+
+    dataset = dataset.map(_remap_fn)
+
+    # Reshape distribution.
+    dataset = dataset.rejection_resample(
+        class_func=lambda x: x, target_dist=target_dist, initial_dist=init_dist)
+
+    get_next = self.getNext(dataset)
+
+    returned = []
+    with self.assertRaises(errors.OutOfRangeError):
+      while True:
+        returned.append(self.evaluate(get_next()))
+
+    classes, _ = zip(*returned)
+    bincount = np.bincount(
+        np.array(classes), minlength=num_classes).astype(
+            np.float32) / len(classes)
+
+    self.assertAllClose(target_dist, bincount, atol=1e-2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExhaustion(self):
+    init_dist = [0.5, 0.5]
+    target_dist = [0.9, 0.1]
+    dataset = dataset_ops.Dataset.range(10000)
+    dataset = dataset.rejection_resample(
+        class_func=lambda x: x % 2,
+        target_dist=target_dist,
+        initial_dist=init_dist)
+
+    get_next = self.getNext(dataset, requires_initialization=True)
+    returned = []
+    with self.assertRaises(errors.OutOfRangeError):
+      while True:
+        returned.append(self.evaluate(get_next()))
+
+    classes, _ = zip(*returned)
+    bincount = np.bincount(
+        np.array(classes), minlength=len(init_dist)).astype(
+            np.float32) / len(classes)
+
+    self.assertAllClose(target_dist, bincount, atol=1e-2)
+
+  @parameterized.parameters(
+      ("float32", "float64"),
+      ("float64", "float32"),
+      ("float64", "float64"),
+      ("float64", None),
+  )
+  def testOtherDtypes(self, target_dtype, init_dtype):
+    target_dist = np.array([0.5, 0.5], dtype=target_dtype)
+
+    if init_dtype is None:
+      init_dist = None
+    else:
+      init_dist = np.array([0.5, 0.5], dtype=init_dtype)
+
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.rejection_resample(
+        class_func=lambda x: x % 2,
+        target_dist=target_dist,
+        initial_dist=init_dist)
+    get_next = self.getNext(dataset, requires_initialization=True)
+    self.evaluate(get_next())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/repeat_test.py

@@ -0,0 +1,316 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.repeat()`."""
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class RepeatTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(count=[0, 3, 7])))
+  def testFiniteRepeat(self, count):
+    """Test a dataset that repeats its input multiple times."""
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+    dataset = dataset_ops.Dataset.from_tensors(components).repeat(count)
+    self.assertEqual(
+        [c.shape for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    self.assertDatasetProduces(dataset, [components] * count)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteRepeat(self):
+    # NOTE(mrry): There's not a good way to test that the sequence is infinite.
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+    dataset = dataset_ops.Dataset.from_tensors(components).repeat(-1)
+    self.assertEqual(
+        [c.shape for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    get_next = self.getNext(dataset)
+    for _ in range(17):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component, result_component)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRepeatRepeat(self):
+    """Test the composition of repeat datasets."""
+    components = (np.array(1), np.array([1, 2, 3]), np.array(37.0))
+    inner_count, outer_count = 7, 14
+
+    dataset = dataset_ops.Dataset.from_tensors(components).repeat(
+        inner_count).repeat(outer_count)
+    self.assertEqual(
+        [c.shape for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    self.assertDatasetProduces(dataset,
+                               [components] * (inner_count * outer_count))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).repeat(1, name="repeat")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class RepeatDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                  parameterized.TestCase):
+
+  def _build_repeat_dataset(self,
+                            num_elements,
+                            num_epochs,
+                            num_outputs=None,
+                            options=None):
+    dataset = dataset_ops.Dataset.range(num_elements).repeat(num_epochs)
+    if num_outputs:
+      range_dataset = dataset_ops.Dataset.range(num_outputs)
+      dataset = dataset_ops.Dataset.zip((dataset, range_dataset))
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def testFiniteRepeat(self, verify_fn, symbolic_checkpoint):
+    num_elements = 10
+    num_epochs = 10
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_repeat_dataset(
+            num_elements, num_epochs, options=options),
+        num_outputs=(num_elements * num_epochs))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def testEmptyRepeat(self, verify_fn, symbolic_checkpoint):
+    num_elements = 10
+    num_epochs = 0
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_repeat_dataset(
+            num_elements, num_epochs, options=options),
+        num_outputs=0)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def testInfiniteRepeat(self, verify_fn, symbolic_checkpoint):
+    num_elements = 10
+    num_epochs = -1
+    num_outputs = 100
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_repeat_dataset(
+            num_elements, num_epochs, num_outputs=num_outputs, options=options),
+        num_outputs=num_outputs)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def testInfiniteEmptyRepeat(self, verify_fn, symbolic_checkpoint):
+    num_elements = 0
+    num_epochs = -1
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_repeat_dataset(
+            num_elements, num_epochs, options=options),
+        num_outputs=0)
+
+
+class RepeatRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 6, 7])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]).repeat(2)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 0])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).repeat(2)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(elements=[0, 5, 10],
+                                              count=[0, 3, 8])))
+  def testFiniteRepeat(self, elements, count):
+    dataset = dataset_ops.Dataset.range(elements).repeat(count)
+    expected_dataset = np.tile(
+        np.arange(
+            start=0, stop=elements, step=1, dtype=dtypes.int64.as_numpy_dtype),
+        count)
+    for i in range(elements * count):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)),
+          expected_dataset[i])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              elements=[0, 3, 5], count_1=[0, 1, 2], count_2=[3, 4, 5])))
+  def testRepeatRepeat(self, elements, count_1, count_2):
+    dataset = dataset_ops.Dataset.range(elements).repeat(count_1).repeat(
+        count_2)
+    expected_dataset = np.tile(
+        np.arange(
+            start=0, stop=elements, step=1, dtype=dtypes.int64.as_numpy_dtype),
+        count_1 * count_2)
+    for i in range(elements * count_1 * count_2):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)),
+          expected_dataset[i])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(elements=[3, 5], count=[None, -1, -2])))
+  def testInfiniteRepeat(self, elements, count):
+    dataset = dataset_ops.Dataset.range(elements).repeat(count=count)
+
+    # Datasets with infinite cardinality do not support random access.
+    with self.assertRaises(errors.FailedPreconditionError):
+      self.evaluate(random_access.at(dataset, index=0))
+
+
+class RepeatGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[41],
+              repetitions=[1, 27],
+              seed=[None, 19],
+              reshuffle_each_iteration=[True, False])))
+  def testRepeat(
+      self,
+      dataset_range: int,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.repeat(repetitions)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(dataset_range)) * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+
+    output_per_iteration = [
+        dataset_output[i : i + dataset_range]
+        for i in range(0, len(dataset_output), dataset_range)]
+    # All sub-ranges should be shuffled.
+    for i in range(1, repetitions):
+      self.assertNotEqual(output_per_iteration[i], list(range(dataset_range)))
+      self.assertNotEqual(output_per_iteration[i], output_per_iteration[i - 1])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(repetitions=[None, 0])))
+  def testInvalidDataset(self, repetitions: Optional[int]):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.repeat(repetitions)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        r"`repeat.*` does not support random access of tf.data datasets."):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class RepeatGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[41],
+              repetitions=[1, 27],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testRepeat(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.repeat(repetitions)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=dataset_range * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/sample_from_datasets_test.py

@@ -0,0 +1,328 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.sample_from_dataset()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import compat as tf_compat
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.platform import test
+
+
+def _weights_type_combinations():
+  return combinations.combine(weights_type=["list", "tensor", "dataset"])
+
+
+def _get_weights_of_type(weights_list, weights_type):
+  if weights_type == "list":
+    return weights_list
+  if weights_type == "tensor":
+    return ops.convert_to_tensor(weights_list, name="weights")
+  return dataset_ops.Dataset.from_tensors(weights_list).repeat()
+
+
+class SampleFromDatasetsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _normalize(self, vec):
+    return vec / vec.sum()
+
+  def _chi2(self, expected, actual):
+    actual = np.asarray(actual)
+    expected = np.asarray(expected)
+    diff = actual - expected
+    chi2 = np.sum(diff * diff / expected, axis=0)
+    return chi2
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _weights_type_combinations()))
+  def testSampleFromDatasets(self, weights_type):
+    random_seed.set_random_seed(1619)
+    num_samples = 5000
+    rand_probs = self._normalize(np.random.random_sample((5,)))
+
+    # Use chi-squared test to assert that the observed distribution matches the
+    # expected distribution. Based on the implementation in
+    # "third_party/tensorflow/python/kernel_tests/multinomial_op_test.py".
+    for probs in [[.85, .05, .1], rand_probs, [1.]]:
+      weights = _get_weights_of_type(np.asarray(probs), weights_type)
+      classes = len(probs)
+
+      # Create a dataset that samples each integer in `[0, num_datasets)`
+      # with probability given by `weights[i]`.
+      dataset = dataset_ops.Dataset.sample_from_datasets([
+          dataset_ops.Dataset.from_tensors(i).repeat() for i in range(classes)
+      ], weights)
+      dataset = dataset.take(num_samples)
+
+      next_element = self.getNext(dataset, requires_initialization=True)
+      freqs = np.zeros([classes])
+      for _ in range(num_samples):
+        freqs[self.evaluate(next_element())] += 1
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element())
+
+      self.assertLess(self._chi2(probs, freqs / num_samples), 1e-2)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _weights_type_combinations()))
+  def testSampleFromDatasetsStoppingOnEmptyDataset(self, weights_type):
+    # Sampling stops when the first dataset is exhausted.
+    weights = _get_weights_of_type(np.asarray([.5, .1, .4]), weights_type)
+    datasets = [
+        dataset_ops.Dataset.from_tensors(np.int64(-1)),
+        dataset_ops.Dataset.from_tensors(np.int64(1)).repeat(),
+        dataset_ops.Dataset.range(10).repeat()
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=True)
+
+    samples_list = self.getIteratorOutput(self.getNext(
+        sample_dataset, requires_initialization=True))
+    self.assertEqual(samples_list.count(-1), 1)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _weights_type_combinations()))
+  def testSampleFromDatasetsSkippingEmptyDataset(self, weights_type):
+    # Sampling skips the first dataset after it becomes empty.
+    weights = _get_weights_of_type(np.asarray([.5, .1, .4]), weights_type)
+    datasets = [
+        dataset_ops.Dataset.from_tensors(np.int64(-1)),
+        dataset_ops.Dataset.from_tensors(np.int64(1)).repeat(),
+        dataset_ops.Dataset.range(10).repeat()
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=False).take(100)
+
+    samples_list = self.getIteratorOutput(self.getNext(
+        sample_dataset, requires_initialization=True))
+    self.assertLen(samples_list, 100)
+    self.assertEqual(samples_list.count(-1), 1)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _weights_type_combinations()))
+  def testSampleFromDatasetsWithZeroWeight(self, weights_type):
+    # Sampling stops when the second dataset is exhausted.
+    weights = _get_weights_of_type(np.asarray([0., 1.]), weights_type)
+    datasets = [
+        dataset_ops.Dataset.from_tensors(-1).repeat(2),
+        dataset_ops.Dataset.from_tensors(1).repeat(2)
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=True)
+    self.assertDatasetProduces(sample_dataset, [1, 1],
+                               requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _weights_type_combinations()))
+  def testSampleFromEmptyDataset(self, weights_type):
+    weights = _get_weights_of_type(np.asarray([1., 0.]), weights_type)
+    datasets = [
+        dataset_ops.Dataset.range(0),
+        dataset_ops.Dataset.range(1).repeat()
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=True)
+    self.assertDatasetProduces(sample_dataset, [],
+                               requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSampleFromDatasetsSkippingDatasetsWithZeroWeight(self):
+    # Sampling skips the first dataset.
+    weights = np.asarray([0., 1.])
+    datasets = [
+        dataset_ops.Dataset.from_tensors(-1).repeat(),
+        dataset_ops.Dataset.from_tensors(1)
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=False)
+    self.assertDatasetProduces(sample_dataset, [1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSampleFromDatasetsAllWeightsAreZero(self):
+    # Sampling skips both datasets.
+    weights = np.asarray([0., 0.])
+    datasets = [
+        dataset_ops.Dataset.from_tensors(-1).repeat(),
+        dataset_ops.Dataset.from_tensors(1).repeat()
+    ]
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=weights, stop_on_empty_dataset=False)
+    self.assertDatasetProduces(sample_dataset, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSampleFromDatasetsCardinality(self):
+    ds1 = dataset_ops.Dataset.from_tensors([1.0]).repeat()
+    ds2 = dataset_ops.Dataset.from_tensors([2.0]).repeat()
+    ds = dataset_ops.Dataset.sample_from_datasets([ds1, ds2])
+    self.assertEqual(self.evaluate(ds.cardinality()), dataset_ops.INFINITE)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSampleFromDatasetsNested(self):
+    ds1 = dataset_ops.Dataset.range(10).window(2)
+    ds2 = dataset_ops.Dataset.range(10, 20).window(2)
+    ds = dataset_ops.Dataset.sample_from_datasets([ds1, ds2],
+                                                  weights=[0.3, 0.7])
+    ds = ds.flat_map(lambda x: x)
+    next_element = self.getNext(ds, requires_initialization=True)
+    self.evaluate(next_element())
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testSampleFromDatasetsRerandomizeEachIterationEpochs(self, rerandomize):
+    if rerandomize is not None and not tf_compat.forward_compatible(
+        2022, 12, 17):
+      self.skipTest(
+          "target functionality not available due to forward compatibility")
+    dataset1 = dataset_ops.Dataset.range(0, 10)
+    dataset2 = dataset_ops.Dataset.range(100, 110)
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        [dataset1, dataset2],
+        seed=42,
+        weights=[0.5, 0.5],
+        stop_on_empty_dataset=True,
+        rerandomize_each_iteration=rerandomize)
+    first_epoch = self.getDatasetOutput(sample_dataset)
+    second_epoch = self.getDatasetOutput(sample_dataset)
+
+    if rerandomize:
+      self.assertNotEqual(first_epoch, second_epoch)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testSampleFromDatasetsRerandomizeRepeatEpochs(self, rerandomize):
+    if rerandomize is not None and not tf_compat.forward_compatible(
+        2022, 12, 17):
+      self.skipTest(
+          "target functionality not available due to forward compatibility")
+    dataset1 = dataset_ops.Dataset.range(0, 10)
+    dataset2 = dataset_ops.Dataset.range(100, 110)
+    sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+        [dataset1, dataset2],
+        seed=42,
+        weights=[0.5, 0.5],
+        stop_on_empty_dataset=True,
+        rerandomize_each_iteration=rerandomize)
+    sample_dataset = sample_dataset.repeat(2)
+    epochs = self.getDatasetOutput(sample_dataset, requires_initialization=True)
+    first_epoch = epochs[:len(epochs) // 2]
+    second_epoch = epochs[len(epochs) // 2:]
+
+    if rerandomize:
+      self.assertNotEqual(first_epoch, second_epoch)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(
+      combinations.times(test_base.v2_eager_only_combinations(),
+                         combinations.combine(rerandomize=[None, True, False])))
+  def testSampleFromDatasetsRerandomizeInsideFunction(self, rerandomize):
+    if rerandomize is not None and not tf_compat.forward_compatible(
+        2022, 12, 17):
+      self.skipTest(
+          "target functionality not available due to forward compatibility")
+    @def_function.function
+    def make_dataset():
+      dataset1 = dataset_ops.Dataset.range(0, 10)
+      dataset2 = dataset_ops.Dataset.range(100, 110)
+      sample_dataset = dataset_ops.Dataset.sample_from_datasets(
+          [dataset1, dataset2],
+          seed=42,
+          weights=[0.5, 0.5],
+          stop_on_empty_dataset=True,
+          rerandomize_each_iteration=rerandomize)
+      return sample_dataset
+
+    sample_dataset = make_dataset()
+    first_epoch = self.getDatasetOutput(sample_dataset)
+    second_epoch = self.getDatasetOutput(sample_dataset)
+
+    if rerandomize:
+      self.assertNotEqual(first_epoch, second_epoch)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrors(self):
+    with self.assertRaisesRegex(ValueError, r"should have the same length"):
+      dataset_ops.Dataset.sample_from_datasets(
+          [dataset_ops.Dataset.range(10),
+           dataset_ops.Dataset.range(20)],
+          weights=[0.25, 0.25, 0.25, 0.25])
+
+    with self.assertRaisesRegex(TypeError, "`tf.float32` or `tf.float64`"):
+      dataset_ops.Dataset.sample_from_datasets(
+          [dataset_ops.Dataset.range(10),
+           dataset_ops.Dataset.range(20)],
+          weights=[1, 1])
+
+    with self.assertRaisesRegex(TypeError, "must have compatible"):
+      dataset_ops.Dataset.sample_from_datasets([
+          dataset_ops.Dataset.from_tensors(0),
+          dataset_ops.Dataset.from_tensors(0.0)
+      ])
+
+    with self.assertRaisesRegex(
+        ValueError, r"Invalid `datasets`. `datasets` should not be empty."):
+      dataset_ops.Dataset.sample_from_datasets(datasets=[], weights=[])
+
+
+class SampleFromDatasetsCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                       parameterized.TestCase):
+
+  def _build_dataset(self, probs, num_samples, options=None):
+    datasets = [
+        dataset_ops.Dataset.from_tensors(i).repeat(None)
+        for i in range(len(probs))
+    ]
+    dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, probs, seed=1813)
+    dataset = dataset.take(num_samples)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_dataset([0.5, 0.5], 100, options),
+        num_outputs=100)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/scan_test.py

@@ -0,0 +1,344 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import itertools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import scan_op
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_v2_toggles
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class ScanTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _counting_dataset(self, start, scan_fn):
+    return dataset_ops.Dataset.from_tensors(0).repeat().scan(
+        initial_state=start, scan_func=scan_fn)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCount(self):
+
+    def make_scan_fn(step):
+      return lambda state, _: (state + step, state)
+
+    def dataset_fn(start, step, take):
+      return self._counting_dataset(start, make_scan_fn(step)).take(take)
+
+    for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
+                                          (10, 2, 10), (10, -1, 10),
+                                          (10, -2, 10)]:
+      next_element = self.getNext(dataset_fn(start_val, step_val, take_val))
+      for expected, _ in zip(
+          itertools.count(start_val, step_val), range(take_val)):
+        self.assertEqual(expected, self.evaluate(next_element()))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFibonacci(self):
+    data = dataset_ops.Dataset.from_tensors(1).repeat(None).scan(
+        initial_state=[0, 1],
+        scan_func=lambda a, _: ([a[1], a[0] + a[1]], a[1]))
+    next_element = self.getNext(data)
+
+    self.assertEqual(1, self.evaluate(next_element()))
+    self.assertEqual(1, self.evaluate(next_element()))
+    self.assertEqual(2, self.evaluate(next_element()))
+    self.assertEqual(3, self.evaluate(next_element()))
+    self.assertEqual(5, self.evaluate(next_element()))
+    self.assertEqual(8, self.evaluate(next_element()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparseCount(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1])),
+          dense_shape=np.array([1, 1]))
+
+    def make_scan_fn(step):
+      return lambda state, _: (_sparse(state.values[0] + step), state)
+
+    def dataset_fn(start, step, take):
+      return self._counting_dataset(_sparse(start),
+                                    make_scan_fn(step)).take(take)
+
+    for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
+                                          (10, 2, 10), (10, -1, 10),
+                                          (10, -2, 10)]:
+      next_element = self.getNext(dataset_fn(start_val, step_val, take_val))
+      for expected, _ in zip(
+          itertools.count(start_val, step_val), range(take_val)):
+        self.assertEqual(expected, self.evaluate(next_element()).values[0])
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorArraySimple(self):
+
+    def scan_fn(ta, x):
+      return (ta.write(ta.size(), x), ta.stack())
+
+    start = tensor_array_ops.TensorArray(
+        size=0, element_shape=[], dtype=dtypes.int64, dynamic_size=True)
+    start = start.write(0, -1)
+
+    ds = dataset_ops.Dataset.range(5).scan(
+        initial_state=start, scan_func=scan_fn)
+
+    self.assertDatasetProduces(
+        ds,
+        expected_output=[
+            [-1],
+            [-1, 0],
+            [-1, 0, 1],
+            [-1, 0, 1, 2],
+            [-1, 0, 1, 2, 3],
+        ],
+        requires_initialization=True,
+        num_test_iterations=2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorArrayWithCondReset(self):
+
+    def empty():
+      return tensor_array_ops.TensorArray(
+          size=0, element_shape=[], dtype=dtypes.int64, dynamic_size=True)
+
+    def scan_fn(ta, x):
+      updated = ta.write(ta.size(), x)
+      next_iter = cond.cond(
+          math_ops.equal(x % 3, 0), empty, lambda: updated)
+      return (next_iter, updated.stack())
+
+    start = empty()
+    start = start.write(0, -1)
+
+    ds = dataset_ops.Dataset.range(6).scan(
+        initial_state=start, scan_func=scan_fn)
+
+    self.assertDatasetProduces(
+        ds,
+        expected_output=[
+            [-1, 0],
+            [1],
+            [1, 2],
+            [1, 2, 3],
+            [4],
+            [4, 5],
+        ],
+        requires_initialization=True,
+        num_test_iterations=2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorArrayWithCondResetByExternalCaptureBreaks(self):
+
+    if control_flow_v2_toggles.control_flow_v2_enabled():
+      self.skipTest("v1 only test")
+
+    empty_ta = tensor_array_ops.TensorArray(
+        size=0, element_shape=[], dtype=dtypes.int64, dynamic_size=True)
+
+    def scan_fn(ta, x):
+      updated = ta.write(ta.size(), x)
+      # Here, capture empty_ta from outside the function.  However, it may be
+      # either a TF1-style TensorArray or an Eager-style TensorArray.
+      next_iter = cond.cond(
+          math_ops.equal(x % 3, 0), lambda: empty_ta, lambda: updated)
+      return (next_iter, updated.stack())
+
+    start = empty_ta
+    start = start.write(0, -1)
+
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        r"construct a new TensorArray inside the function"):
+      dataset_ops.Dataset.range(6).scan(initial_state=start, scan_func=scan_fn)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChangingStateShape(self):
+    # Test the fixed-point shape invariant calculations: start with
+    # initial values with known shapes, and use a scan function that
+    # changes the size of the state on each element.
+    def _scan_fn(state, input_value):
+      # Statically known rank, but dynamic length.
+      ret_longer_vector = array_ops.concat([state[0], state[0]], 0)
+      # Statically unknown rank.
+      ret_larger_rank = array_ops.expand_dims(state[1], 0)
+      return (ret_longer_vector, ret_larger_rank), (state, input_value)
+
+    dataset = dataset_ops.Dataset.from_tensors(0).repeat(5).scan(
+        initial_state=([0], 1), scan_func=_scan_fn)
+    self.assertEqual(
+        [None],
+        dataset_ops.get_legacy_output_shapes(dataset)[0][0].as_list())
+    self.assertIs(None,
+                  dataset_ops.get_legacy_output_shapes(dataset)[0][1].ndims)
+    self.assertEqual([],
+                     dataset_ops.get_legacy_output_shapes(dataset)[1].as_list())
+
+    next_element = self.getNext(dataset)
+
+    for i in range(5):
+      (longer_vector_val, larger_rank_val), _ = self.evaluate(next_element())
+      self.assertAllEqual([0] * (2**i), longer_vector_val)
+      self.assertAllEqual(np.array(1, ndmin=i), larger_rank_val)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIncorrectStateType(self):
+
+    def _scan_fn(state, _):
+      return constant_op.constant(1, dtype=dtypes.int64), state
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        TypeError,
+        "The element types for the new state must match the initial state."):
+      dataset.scan(
+          initial_state=constant_op.constant(1, dtype=dtypes.int32),
+          scan_func=_scan_fn)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStateLengthMismatch(self):
+
+    def _scan_fn(state, _):
+      return (state, state), state
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        TypeError,
+        "The state returned by `scan_func` must have the same number of "
+        "elements as the initial state.",
+    ):
+      dataset.scan(
+          initial_state=constant_op.constant(1, dtype=dtypes.int32),
+          scan_func=_scan_fn,
+      )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIncorrectReturnType(self):
+
+    def _scan_fn(unused_state, unused_input_value):
+      return constant_op.constant(1, dtype=dtypes.int64)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        TypeError,
+        "`scan_func` should return a pair consisting of new state and the "
+        "output value."):
+      dataset.scan(
+          initial_state=constant_op.constant(1, dtype=dtypes.int32),
+          scan_func=_scan_fn)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPreserveCardinality(self):
+
+    def scan_fn(state, val):
+
+      def py_fn(_):
+        raise StopIteration()
+
+      return state, script_ops.py_func(py_fn, [val], dtypes.int64)
+
+    dataset = dataset_ops.Dataset.from_tensors(0).scan(
+        initial_state=constant_op.constant(1), scan_func=scan_fn)
+    get_next = self.getNext(dataset)
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.combine(
+          tf_api_version=2, mode="eager", use_default_device=[True, False]))
+  def testUseDefaultDevice(self, use_default_device):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPUs available.")
+
+    weights = variables.Variable(initial_value=array_ops.zeros((1000, 1000)))
+    result = variables.Variable(initial_value=array_ops.zeros((1000, 1000)))
+
+    def scan_fn(state, sample):
+      product = math_ops.matmul(sample, weights)
+      result.assign_add(product)
+      with ops.colocate_with(product):
+        device = test_ops.device_placement_op()
+      return state, device
+
+    data = variables.Variable(initial_value=array_ops.zeros((1, 1000, 1000)))
+    dataset = dataset_ops.Dataset.from_tensor_slices(data)
+    dataset = scan_op._ScanDataset(
+        dataset, np.int64(1), scan_fn, use_default_device=use_default_device)
+    get_next = self.getNext(dataset)
+
+    if use_default_device:
+      self.assertIn(b"CPU:0", self.evaluate(get_next()))
+    else:
+      self.assertIn(b"GPU:0", self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).scan(
+        0, lambda x, y: (y, y), name="scan")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class ScanCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                         parameterized.TestCase):
+
+  def _build_dataset(self, num_elements, symbolic_checkpoint):
+    dataset = dataset_ops.Dataset.from_tensors(1).repeat(num_elements)
+    dataset = dataset.scan(
+        initial_state=[0, 1],
+        scan_func=lambda a, _: ([a[1], a[0] + a[1]], a[1]))
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    return dataset.with_options(options)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    num_outputs = 10
+    verify_fn(
+        self,
+        lambda: self._build_dataset(num_outputs, symbolic_checkpoint),
+        num_outputs=num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/shard_test.py

@@ -0,0 +1,266 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.shard()`."""
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class ShardTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimpleCase(self):
+    dataset = dataset_ops.Dataset.range(10).shard(5, 2)
+    self.assertDatasetProduces(dataset, expected_output=[2, 7])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedData(self):
+    dataset_a = dataset_ops.Dataset.range(10)
+    dataset_b = dataset_ops.Dataset.range(10, 0, -1)
+    dataset = dataset_ops.Dataset.zip((dataset_a, dataset_b)).shard(5, 2)
+    self.assertDatasetProduces(dataset, expected_output=[(2, 8), (7, 3)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOffsetZero(self):
+    dataset = dataset_ops.Dataset.range(10).shard(5, 0)
+    self.assertDatasetProduces(dataset, expected_output=[0, 5])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOffsetGreaterNumShards(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(10).shard(5, 7)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNegativeOffset(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(10).shard(5, -3)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNegativeNumShards(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(10).shard(-3, 1)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZeroNumShards(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.range(10).shard(0, 1)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIteratorEndsBeforeFirstElem(self):
+    dataset = dataset_ops.Dataset.range(1).shard(5, 2)
+    self.assertDatasetProduces(dataset, expected_output=[])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLargerWorkerPool(self):
+    dataset = dataset_ops.Dataset.range(10).shard(7, 5)
+    self.assertDatasetProduces(dataset, expected_output=[5])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIndexEqualsNumShards(self):
+    dataset = dataset_ops.Dataset.range(10).shard(5, 4)
+    self.assertDatasetProduces(dataset, expected_output=[4, 9])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIndexEqualsNumShards2(self):
+    dataset = dataset_ops.Dataset.range(10).shard(4, 3)
+    self.assertDatasetProduces(dataset, expected_output=[3, 7])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNumShardsLargerThanDataset(self):
+    dataset = dataset_ops.Dataset.range(10).shard(20, 5)
+    self.assertDatasetProduces(dataset, expected_output=[5])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).shard(1, 0, name="shard")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class ShardCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                          parameterized.TestCase):
+
+  def _build_dataset(self, num_elements, num_shards, index, options=None):
+    dataset = dataset_ops.Dataset.range(num_elements).shard(num_shards, index)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True]),
+          combinations.combine(
+              elems=[10, 100], num_shards=[2, 5], index=[0, 1])))
+  def test(self, verify_fn, symbolic_checkpoint, elems, num_shards, index):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_dataset(elems, num_shards, index, options),
+        num_outputs=elems // num_shards)
+
+
+class ShardRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 2, 3, 4])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(4).shard(num_shards=2, index=0)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).shard(
+        num_shards=2, index=1)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=0))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNumShardsAndIndexLessThanNumElements(self):
+    dataset = dataset_ops.Dataset.range(10).shard(5, 0)
+    self.assertEqual(0, self.evaluate(random_access.at(dataset, 0)))
+    self.assertEqual(5, self.evaluate(random_access.at(dataset, 1)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNumShardsGreaterThanNumElementsIndexLess(self):
+    dataset = dataset_ops.Dataset.range(7).shard(8, 3)
+    self.assertEqual(3, self.evaluate(random_access.at(dataset, 0)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNumShardsAndIndexGreaterThanNumElements(self):
+    dataset = dataset_ops.Dataset.range(13).shard(23, 21)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 0))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              elements=[0, 10, 50],
+              num_shards=[5, 7, 10],
+              index=[0, 1, 2, 3, 4],
+          )))
+  def testMultipleCombinations(self, elements, num_shards, index):
+    components = range(elements)
+    dataset = dataset_ops.Dataset.range(elements).shard(
+        num_shards=num_shards, index=index)
+    len_dataset = self.evaluate(dataset.cardinality())
+    for i in range(self.evaluate(dataset.cardinality())):
+      self.assertAllEqual(components[index + (num_shards * i)],
+                          self.evaluate(random_access.at(dataset, i)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=len_dataset))
+
+
+class ShardGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              num_shards=[1, 3, 5],
+              shard_index=[0, 1, 2, 4],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def testShard(
+      self,
+      dataset_range: int,
+      num_shards: int,
+      shard_index: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    if shard_index >= num_shards:
+      return
+
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.shard(num_shards, shard_index)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(shard_index, dataset_range, num_shards))
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+
+
+class ShardGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              num_shards=[1, 3],
+              shard_index=[0, 1, 2],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testShard(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      num_shards: int,
+      shard_index: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    if shard_index >= num_shards:
+      return
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.shard(num_shards, shard_index)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=len(range(shard_index, dataset_range, num_shards)),
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/shuffle_test.py

@@ -0,0 +1,708 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.shuffle()`."""
+import collections
+import functools
+import sys
+
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python import pywrap_sanitizers
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+
+
+def make_variable_size_dataset(per_epoch_data):
+  repeat_counter = [0]
+
+  def gen():
+    for each in per_epoch_data[repeat_counter[0]]:
+      yield each
+    repeat_counter[0] += 1
+
+  return dataset_ops.Dataset.from_generator(gen, dtypes.int64)
+
+
+class ShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    components = (
+        np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]),
+        np.array([9.0, 10.0, 11.0, 12.0])
+    )
+
+    def dataset_fn(count=5, buffer_size=None, seed=0):
+      repeat_dataset = (
+          dataset_ops.Dataset.from_tensor_slices(components).repeat(count))
+      if buffer_size:
+        shuffle_dataset = repeat_dataset.shuffle(buffer_size, seed)
+
+        self.assertEqual(
+            tuple([c.shape[1:] for c in components]),
+            dataset_ops.get_legacy_output_shapes(shuffle_dataset))
+        return shuffle_dataset
+      else:
+        return repeat_dataset
+
+    # First run without shuffling to collect the "ground truth".
+    get_next = self.getNext(dataset_fn())
+    unshuffled_elements = []
+    for _ in range(20):
+      unshuffled_elements.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Assert that the shuffled dataset has the same elements as the
+    # "ground truth".
+    get_next = self.getNext(dataset_fn(buffer_size=100, seed=37))
+    shuffled_elements = []
+    for _ in range(20):
+      shuffled_elements.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertAllEqual(sorted(unshuffled_elements), sorted(shuffled_elements))
+
+    # Assert that shuffling twice with the same seeds gives the same sequence.
+    get_next = self.getNext(dataset_fn(buffer_size=100, seed=37))
+    reshuffled_elements_same_seed = []
+    for _ in range(20):
+      reshuffled_elements_same_seed.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertEqual(shuffled_elements, reshuffled_elements_same_seed)
+
+    # Assert that shuffling twice with a different seed gives a different
+    # permutation of the same elements.
+    get_next = self.getNext(dataset_fn(
+        buffer_size=100, seed=constant_op.constant(137, dtype=dtypes.int64)))
+    reshuffled_elements_different_seed = []
+    for _ in range(20):
+      reshuffled_elements_different_seed.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertNotEqual(shuffled_elements, reshuffled_elements_different_seed)
+    self.assertAllEqual(
+        sorted(shuffled_elements), sorted(reshuffled_elements_different_seed))
+
+    # Assert that the shuffled dataset has the same elements as the
+    # "ground truth" when the buffer size is smaller than the input
+    # dataset.
+    get_next = self.getNext(dataset_fn(buffer_size=2, seed=37))
+    reshuffled_elements_small_buffer = []
+    for _ in range(20):
+      reshuffled_elements_small_buffer.append(self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    self.assertAllEqual(
+        sorted(unshuffled_elements), sorted(reshuffled_elements_small_buffer))
+
+    # Test the case of shuffling an empty dataset.
+    get_next = self.getNext(dataset_fn(count=0, buffer_size=100, seed=37))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testSeedZero(self):
+    """Test for same behavior when the seed is a Python or Tensor zero."""
+    iterator = dataset_ops.make_one_shot_iterator(
+        dataset_ops.Dataset.range(10).shuffle(10, seed=0))
+    get_next = iterator.get_next()
+
+    elems = []
+    with self.cached_session() as sess:
+      for _ in range(10):
+        elems.append(sess.run(get_next))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+    seed_placeholder = array_ops.placeholder(dtypes.int64, shape=[])
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.range(10).shuffle(10, seed=seed_placeholder))
+    get_next = iterator.get_next()
+
+    with self.cached_session() as sess:
+      sess.run(iterator.initializer, feed_dict={seed_placeholder: 0})
+      for elem in elems:
+        self.assertEqual(elem, sess.run(get_next))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDefaultArguments(self):
+    components = [0, 1, 2, 3, 4]
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components).shuffle(5).repeat()
+    )
+    get_next = self.getNext(dataset)
+    counts = collections.defaultdict(lambda: 0)
+    for _ in range(10):
+      for _ in range(5):
+        counts[self.evaluate(get_next())] += 1
+
+    for i in range(5):
+      self.assertEqual(10, counts[i])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              buffer_size=[None, 10, 200],
+              seed=[None, 42],
+              use_tensor_input=[True, False])))
+  def testTensorInput(self, dataset_range, buffer_size, seed, use_tensor_input):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    unshuffled_output = self.getDatasetOutput(dataset)
+
+    if buffer_size:
+      buffer_size = (
+          constant_op.constant(buffer_size, dtype=dtypes.int64)
+          if use_tensor_input else buffer_size)
+    else:
+      buffer_size = dataset.cardinality()
+    seed = (constant_op.constant(seed, dtype=dtypes.int64)
+            if seed and use_tensor_input else seed)
+
+    shuffled_dataset = dataset.shuffle(buffer_size, seed=seed)
+    shuffled_output = self.getDatasetOutput(shuffled_dataset)
+    self.assertEqual(unshuffled_output, list(range(dataset_range)))
+    self.assertCountEqual(shuffled_output, unshuffled_output)
+    self.assertNotEqual(shuffled_output, unshuffled_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownCardinality(self):
+    components = [0, 1, 2, 3, 4]
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).shuffle(
+        dataset_ops.UNKNOWN
+    )
+    get_next = self.getNext(dataset)
+    counts = collections.defaultdict(lambda: 0)
+    for _ in range(1):
+      for _ in range(5):
+        counts[self.evaluate(get_next())] += 1
+
+    for i in range(5):
+      self.assertEqual(1, counts[i])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownCardinalityWithRepeatedShuffle(self):
+    components = [0, 1, 2, 3, 4]
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(components)
+        .shuffle(dataset_ops.UNKNOWN)
+        .repeat()
+    )
+    get_next = self.getNext(dataset)
+    counts = collections.defaultdict(lambda: 0)
+    for _ in range(10):
+      for _ in range(5):
+        counts[self.evaluate(get_next())] += 1
+
+    for i in range(5):
+      self.assertEqual(10, counts[i])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownCardinalityWithIncreasingBufferSize(self):
+    epoch_1 = list(range(5))
+    epoch_2 = list(range(10, 17))
+    epoch_3 = list(range(20, 28))
+
+    ds = make_variable_size_dataset([epoch_1, epoch_2, epoch_3])
+    ds = ds.shuffle(dataset_ops.UNKNOWN).repeat(3)
+
+    expected = epoch_1 + epoch_2 + epoch_3
+    self.assertDatasetProduces(ds, expected, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownCardinalityWithVariableBufferSize(self):
+    epoch_1 = list(range(5))
+    epoch_2 = list(range(10, 13))
+    epoch_3 = list(range(20, 27))
+
+    ds = make_variable_size_dataset([epoch_1, epoch_2, epoch_3])
+    ds = ds.shuffle(dataset_ops.UNKNOWN).repeat(3)
+
+    expected = epoch_1 + epoch_2 + epoch_3
+    self.assertDatasetProduces(ds, expected, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInputInitializations(self):
+    num_rounds = 3
+    def compute_orders(dataset):
+      orders = []
+      for _ in range(num_rounds):
+        orders.append(self.getDatasetOutput(dataset))
+      return orders
+
+    dataset = dataset_ops.Dataset.range(10).shuffle(10, seed=1)
+    first_orders = compute_orders(dataset)
+    dataset = dataset_ops.Dataset.range(10)
+
+    # Adding shuffle(1) should not change the order.
+    dataset = dataset_ops.Dataset.range(10).shuffle(10, seed=1).shuffle(1)
+    second_orders = compute_orders(dataset)
+    self.assertEqual(first_orders, second_orders)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.graph_only_combinations(),
+          combinations.combine(reshuffle=[True, False]),
+          combinations.combine(graph_seed=38, op_seed=None) +
+          combinations.combine(graph_seed=None, op_seed=42) +
+          combinations.combine(graph_seed=38, op_seed=42)))
+  def testShuffleSeed(self, reshuffle, graph_seed, op_seed):
+    results = []
+    for _ in range(2):
+      with ops.Graph().as_default() as g:
+        random_seed.set_random_seed(graph_seed)
+        dataset = dataset_ops.Dataset.range(10).shuffle(
+            10, seed=op_seed, reshuffle_each_iteration=reshuffle).repeat(3)
+        iterator = dataset_ops.make_one_shot_iterator(dataset)
+        next_element = iterator.get_next()
+
+        run_results = []
+        with self.session(graph=g) as sess:
+          for _ in range(30):
+            run_results.append(sess.run(next_element))
+          with self.assertRaises(errors.OutOfRangeError):
+            sess.run(next_element)
+        results.append(run_results)
+
+    self.assertAllEqual(results[0], results[1])
+
+  # TODO(b/117581999): enable this test for eager-mode.
+  @combinations.generate(
+      combinations.times(
+          test_base.graph_only_combinations(),
+          combinations.combine(
+              reshuffle=[True, False], initializable=[True, False])))
+  def testMultipleIterators(self, reshuffle, initializable):
+    with ops.Graph().as_default() as g:
+      dataset = dataset_ops.Dataset.range(100).shuffle(
+          10, reshuffle_each_iteration=reshuffle).repeat(3)
+
+      if initializable:
+        iterators = [dataset_ops.make_initializable_iterator(dataset)
+                     for _ in range(2)]
+      else:
+        iterators = [dataset_ops.make_one_shot_iterator(dataset)
+                     for _ in range(2)]
+
+      results = []
+      with self.session(graph=g) as sess:
+        for iterator in iterators:
+          if initializable:
+            sess.run(iterator.initializer)
+          next_element = iterator.get_next()
+          run_results = []
+          for _ in range(300):
+            run_results.append(sess.run(next_element))
+          with self.assertRaises(errors.OutOfRangeError):
+            sess.run(next_element)
+
+          results.append(run_results)
+
+        self.assertNotEqual(results[0], results[1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleManyEmptyEpochs(self):
+    sizes = [0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 0, 0]
+    sizes_iter = iter(sizes)
+    def gen():
+      for i in range(next(sizes_iter)):
+        yield i
+
+    dataset = dataset_ops.Dataset.from_generator(
+        gen, output_signature=tensor_spec.TensorSpec((), dtypes.int64))
+    dataset = dataset.shuffle(10).repeat(len(sizes)).take(3)
+    self.assertDatasetProduces(dataset, [0, 0, 1], assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleInfiniteRepeatNonemptyFollowedByEmpty(self):
+    sizes = [1, 0, 2, 10]
+    sizes_iter = iter(sizes)
+    def gen():
+      for i in range(next(sizes_iter)):
+        yield i
+
+    dataset = dataset_ops.Dataset.from_generator(
+        gen, output_signature=tensor_spec.TensorSpec((), dtypes.int64))
+    dataset = dataset.shuffle(10).repeat().take(3)
+    self.assertDatasetProduces(dataset, [0, 0, 1], assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(reshuffle=[True, False], seed=[None, 42])))
+  def testReshuffleRepeatEpochs(self, reshuffle, seed):
+    dataset = dataset_ops.Dataset.range(10).shuffle(
+        10, seed=seed, reshuffle_each_iteration=reshuffle).repeat(2)
+    next_element = self.getNext(dataset)
+
+    first_epoch = []
+    for _ in range(10):
+      first_epoch.append(self.evaluate(next_element()))
+
+    second_epoch = []
+    for _ in range(10):
+      second_epoch.append(self.evaluate(next_element()))
+
+    self.assertEqual(first_epoch == second_epoch, not reshuffle)
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(tf_api_version=2, mode="eager"),
+          combinations.combine(reshuffle=[True, False], seed=[None, 42])))
+  def testReshuffleIterationEpochs(self, reshuffle, seed):
+    # TensorFlow unit tests set the global graph seed. We unset it here so that
+    # we can control determinism via the `seed` parameter.
+    random_seed.set_random_seed(None)
+    dataset = dataset_ops.Dataset.range(10).shuffle(
+        10, seed=seed, reshuffle_each_iteration=reshuffle)
+
+    first_epoch = self.getDatasetOutput(dataset)
+    second_epoch = self.getDatasetOutput(dataset)
+
+    self.assertEqual(first_epoch == second_epoch, not reshuffle)
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testShuffleV2ResourceCapture(self):
+
+    def make_dataset():
+      ids = dataset_ops.Dataset.range(10)
+      ids = ids.shuffle(1)
+
+      def interleave_fn(dataset, _):
+        return dataset
+
+      dataset = dataset_ops.Dataset.range(1)
+      dataset = dataset.interleave(functools.partial(interleave_fn, ids))
+      return dataset
+
+    results = []
+    for elem in make_dataset():
+      results.append(elem.numpy())
+
+    self.assertAllEqual(results, range(10))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(reshuffle=[True, False], seed=[None, 42])))
+  def testReshuffleSeparateTransformations(self, reshuffle, seed):
+    dataset = dataset_ops.Dataset.range(10)
+
+    first_epoch = []
+    for elem in dataset.shuffle(
+        10, seed=seed, reshuffle_each_iteration=reshuffle):
+      first_epoch.append(elem.numpy())
+
+    second_epoch = []
+    for elem in dataset.shuffle(
+        10, seed=seed, reshuffle_each_iteration=reshuffle):
+      second_epoch.append(elem.numpy())
+
+    self.assertEqual(first_epoch != second_epoch, seed is None)
+
+  @combinations.generate(combinations.combine(tf_api_version=2, mode="eager"))
+  def testShuffleV2InFunction(self):
+    counter_var = variables.Variable(0)
+
+    @def_function.function
+    def consume():
+      ds = dataset_ops.Dataset.range(10)
+      ds = ds.shuffle(1)
+      for _ in ds:
+        counter_var.assign(counter_var + 1)
+
+    consume()
+    self.assertAllEqual(self.evaluate(counter_var), 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.from_tensors(1)
+
+    def map_fn(x):
+      with ops.control_dependencies([check_ops.assert_equal(x, 0)]):
+        return x
+
+    dataset = dataset.map(map_fn)
+    dataset = dataset.cache()
+    dataset = dataset.shuffle(buffer_size=10).repeat()
+
+    get_next = self.getNext(dataset)
+
+    # First time around, we get an error for the failed assertion.
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+
+    # Second time around, we get an EOF because the cached dataset is empty.
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(reshuffle=[True, False])))
+  def testDontRerandomizeOnReplicate(self, reshuffle):
+    random_seed.set_random_seed(None)
+    # Since the seed generator configuration is preserved across serialization
+    # of the dataset, each instantiation of the shuffle dataset
+    # should preserve the shuffle order if reshuffle=False. To preserve the
+    # shuffle order, the original dataset must be kept alive, since if the
+    # original dataset was destroyed, its seeds would also be destroyed.
+    num_elements = 100
+    dataset_1 = dataset_ops.Dataset.range(num_elements)
+    dataset_2 = dataset_1.shuffle(
+        num_elements, reshuffle_each_iteration=reshuffle)
+
+    shuffle_1 = self.getDatasetOutput(dataset_2)
+    dataset_3 = self.graphRoundTrip(dataset_2, allow_stateful=True)
+    shuffle_2 = self.getDatasetOutput(dataset_3)
+
+    self.assertCountEqual(shuffle_1, shuffle_2)
+    if reshuffle:
+      self.assertNotEqual(shuffle_1, shuffle_2)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCheckpointLargeBuffer(self):
+    if (pywrap_sanitizers.is_asan_enabled() or
+        pywrap_sanitizers.is_tsan_enabled() or
+        pywrap_sanitizers.is_msan_enabled()):
+      self.skipTest("Skip to avoid OOM when using sanitizers.")
+    if sys.platform == "darwin":
+      self.skipTest("Skip to avoid memory issues on mac.")
+    dataset = dataset_ops.Dataset.range(12).batch(2)
+    dataset = dataset.map(
+        # Create tensors of size 512M.
+        lambda seed: stateless_random_ops.stateless_random_uniform(
+            (128, 1024, 1024), seed, dtype=dtypes.float32
+        )
+    )
+    dataset = dataset.shuffle(buffer_size=6)
+    iterator = iter(dataset)
+    next(iterator)  # Request an element to fill the shuffle buffer
+    ckpt = trackable_utils.Checkpoint(iterator=iterator)
+    manager = checkpoint_management.CheckpointManager(
+        ckpt, self.get_temp_dir(), max_to_keep=1)
+    manager.save()
+    del dataset
+    del iterator
+    manager.restore_or_initialize()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).shuffle(1, name="shuffle")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class ShuffleCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                            parameterized.TestCase):
+
+  def _build_shuffle_dataset(
+      self,
+      range_limit=10,
+      num_repeats=5,
+      buffer_size=5,
+      seed=None,
+      reshuffle_each_iteration=None,
+      symbolic_checkpoint=None,
+  ):
+    dataset = (
+        dataset_ops.Dataset.range(range_limit)
+        .shuffle(
+            buffer_size,
+            seed=seed,
+            reshuffle_each_iteration=reshuffle_each_iteration,
+        )
+        .repeat(num_repeats)
+    )
+
+    if symbolic_checkpoint:
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      dataset = dataset.with_options(options)
+
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              symbolic_checkpoint=[True, False],
+              reshuffle_each_iteration=[True, False],
+              buffer_size=[1, 3, 5, 8, 10, dataset_ops.UNKNOWN],
+          ),
+      )
+  )
+  def test(
+      self,
+      verify_fn,
+      symbolic_checkpoint,
+      reshuffle_each_iteration,
+      buffer_size,
+  ):
+    seed = 55
+    range_limit = 5
+    num_repeats = 2
+    num_outputs = range_limit * num_repeats
+    # pylint: disable=g-long-lambda
+    verify_fn(
+        self,
+        lambda: self._build_shuffle_dataset(
+            range_limit=range_limit,
+            num_repeats=num_repeats,
+            buffer_size=buffer_size,
+            seed=seed,
+            reshuffle_each_iteration=reshuffle_each_iteration,
+            symbolic_checkpoint=symbolic_checkpoint,
+        ),
+        num_outputs,
+    )
+
+  @combinations.generate(
+      combinations.combine(
+          tf_api_version=1,
+          mode=["graph"],
+          reshuffle_each_iteration=[True, False],
+          buffer_size=[1, 3, 5, 8, 10]))
+  def testMultipleIterators(self, reshuffle_each_iteration, buffer_size):
+    range_limit = 5
+    num_repeats = 2
+    num_outputs = range_limit * num_repeats
+
+    def ds_fn():
+      # pylint: disable=cell-var-from-loop
+      return self._build_shuffle_dataset(
+          range_limit=range_limit,
+          num_repeats=num_repeats,
+          buffer_size=buffer_size,
+          seed=None,  # Iterator seeds are generated non-deterministically.
+          reshuffle_each_iteration=reshuffle_each_iteration)
+      # pylint: enable=cell-var-from-loop
+
+    with ops.Graph().as_default() as g:
+      ds = ds_fn()
+      iterators = [ds.make_one_shot_iterator(), ds.make_one_shot_iterator()]
+      get_next_ops = [it.get_next() for it in iterators]
+      saveables = [
+          contrib_iterator_ops.make_saveable_from_iterator(it)
+          for it in iterators
+      ]
+      for saveable in saveables:
+        ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable)
+      saver = saver_lib.Saver(allow_empty=True)
+      with self.session(graph=g) as sess:
+        self._save(sess, saver)
+        expected = [self.evaluate(get_next_ops) for _ in range(num_outputs)]
+        self._restore(saver, sess)
+        actual = [self.evaluate(get_next_ops) for _ in range(num_outputs)]
+        self.match(expected, actual)
+
+
+class ShuffleRandomAccessTest(test_base.DatasetTestBase,
+                              parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidIndex(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3
+                                                     ]).shuffle(buffer_size=100)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, -1))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        []).shuffle(buffer_size=100)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, 0))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testBasicWithoutSeedEager(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5])
+    shuffled_dataset = dataset.shuffle(buffer_size=100)
+
+    dataset_array = []
+    shuffled_dataset_array = []
+
+    for i in range(5):
+      shuffled_dataset_array.append(
+          self.evaluate(random_access.at(shuffled_dataset, i)))
+      dataset_array.append(self.evaluate(random_access.at(dataset, i)))
+    self.assertAllEqual(sorted(dataset_array), sorted(shuffled_dataset_array))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSameSeedReturnsSameSequence(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3, 4, 5])
+    shuffled_dataset = dataset.shuffle(buffer_size=100, seed=5)
+    shuffled_dataset_2 = dataset.shuffle(buffer_size=100, seed=5)
+
+    shuffled_dataset_array = []
+    shuffled_dataset_array_2 = []
+
+    for i in range(5):
+      shuffled_dataset_array.append(
+          self.evaluate(random_access.at(shuffled_dataset, i)))
+      shuffled_dataset_array_2.append(
+          self.evaluate(random_access.at(shuffled_dataset_2, i)))
+    self.assertAllEqual(shuffled_dataset_array, shuffled_dataset_array_2)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDifferentSeedDifferentSequence(self):
+    components = list(range(1000))
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    shuffled_dataset = dataset.shuffle(buffer_size=1000, seed=124)
+    shuffled_dataset_2 = dataset.shuffle(buffer_size=1000, seed=51)
+
+    shuffled_dataset_array = []
+    shuffled_dataset_array_2 = []
+
+    for i in range(1000):
+      shuffled_dataset_array.append(
+          self.evaluate(random_access.at(shuffled_dataset, i)))
+      shuffled_dataset_array_2.append(
+          self.evaluate(random_access.at(shuffled_dataset_2, i)))
+    self.assertNotEqual(shuffled_dataset_array, shuffled_dataset_array_2)
+    self.assertAllEqual(
+        sorted(shuffled_dataset_array), sorted(shuffled_dataset_array_2))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/skip_test.py

@@ -0,0 +1,216 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.skip()`."""
+
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class SkipTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(count=[-1, 0, 4, 10, 25])))
+  def testBasic(self, count):
+    components = (np.arange(10),)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).skip(count)
+    self.assertEqual(
+        [c.shape[1:] for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    start_range = min(count, 10) if count != -1 else 10
+    self.assertDatasetProduces(
+        dataset,
+        [tuple(components[0][i:i + 1]) for i in range(start_range, 10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).skip(0, name="skip")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class SkipDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  def _build_skip_dataset(self, count, options=None):
+    dataset = dataset_ops.Dataset.range(100).skip(count)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True]),
+          combinations.combine(count=[50], num_outputs=[50]) +
+          combinations.combine(count=[200, 100, -1], num_outputs=[0]) +
+          combinations.combine(count=[0], num_outputs=[100])))
+  def test(self, verify_fn, count, num_outputs, symbolic_checkpoint):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self._build_skip_dataset(count, options),
+              num_outputs)
+
+
+class SkipRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 2, 3])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(10).skip(8)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 0])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).skip(8)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testBasic(self):
+    dataset = dataset_ops.Dataset.range(11).skip(3)
+    for i in range(8):
+      self.assertEqual(self.evaluate(random_access.at(dataset, index=i)), i + 3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(skip=[-2, -1])))
+  def testNegativeSkip(self, skip):
+    dataset = dataset_ops.Dataset.range(11).skip(skip)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=0))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(skip=[5, 8])))
+  def testSkipGreaterThanNumElements(self, skip):
+    dataset = dataset_ops.Dataset.range(4).skip(skip)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=0))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(elements=[0, 5, 10], skip=[-1, 0, 5, 15])))
+  def testMultipleCombinations(self, elements, skip):
+    dataset = dataset_ops.Dataset.range(elements).skip(skip)
+    for i in range(self.evaluate(dataset.cardinality())):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)), i + skip)
+
+
+class SkipGlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              count=[0, 2],
+              repetitions=[1, 2],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def testSkip(
+      self,
+      dataset_range: int,
+      count: int,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.skip(count)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(count, dataset_range)) * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(dataset_output, self.evaluate(dataset.cardinality()))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(skip=[-2, -1])))
+  def testNegativeSkip(self, skip: int):
+    dataset = dataset_ops.Dataset.range(10).skip(skip)
+    with self.assertRaises(errors.FailedPreconditionError):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class SkipGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              count=[0, 2],
+              repetitions=[1, 2],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testSkip(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      count: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.skip(count)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=(dataset_range - count) * repetitions,
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/sparse_batch_test.py

@@ -0,0 +1,123 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.sparse_batch`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class DenseToSparseBatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    components = np.random.randint(12, size=(100,)).astype(np.int32)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+        lambda x: array_ops.fill([x], x)).sparse_batch(4, [12])
+    get_next = self.getNext(dataset)
+
+    for start in range(0, len(components), 4):
+      results = self.evaluate(get_next())
+      self.assertAllEqual([[i, j]
+                           for i, c in enumerate(components[start:start + 4])
+                           for j in range(c)], results.indices)
+      self.assertAllEqual(
+          [c for c in components[start:start + 4] for _ in range(c)],
+          results.values)
+      self.assertAllEqual([min(4,
+                               len(components) - start), 12],
+                          results.dense_shape)
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithUnknownShape(self):
+    components = np.random.randint(5, size=(40,)).astype(np.int32)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+        lambda x: array_ops.fill([x, x], x)).sparse_batch(4, [5, None])
+
+    get_next = self.getNext(dataset)
+
+    for start in range(0, len(components), 4):
+      results = self.evaluate(get_next())
+      self.assertAllEqual([[i, j, z]
+                           for i, c in enumerate(components[start:start + 4])
+                           for j in range(c)
+                           for z in range(c)], results.indices)
+      self.assertAllEqual([
+          c for c in components[start:start + 4] for _ in range(c)
+          for _ in range(c)
+      ], results.values)
+      self.assertAllEqual([
+          min(4,
+              len(components) - start), 5,
+          np.max(components[start:start + 4])
+      ], results.dense_shape)
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithInvalidShape(self):
+    input_tensor = array_ops.constant([[1]])
+    with self.assertRaisesRegex(ValueError, "Dimension -2 must be >= 0"):
+      dataset_ops.Dataset.from_tensors(input_tensor).sparse_batch(4, [-2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShapeErrors(self):
+
+    def dataset_fn(input_tensor):
+      return dataset_ops.Dataset.from_tensors(input_tensor).sparse_batch(
+          4, [12])
+
+    # Initialize with an input tensor of incompatible rank.
+    get_next = self.getNext(dataset_fn([[1]]))
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "incompatible with the row shape"):
+      self.evaluate(get_next())
+
+    # Initialize with an input tensor that is larger than `row_shape`.
+    get_next = self.getNext(dataset_fn(np.int32(range(13))))
+    with self.assertRaisesRegex(errors.DataLossError,
+                                "larger than the row shape"):
+      self.evaluate(get_next())
+
+
+class DenseToSparseBatchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                       parameterized.TestCase):
+
+  def _build_dataset(self, components):
+    return dataset_ops.Dataset.from_tensor_slices(components).map(
+        lambda x: array_ops.fill([x], x)).sparse_batch(4, [12])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    components = np.random.randint(5, size=(40,)).astype(np.int32)
+
+    num_outputs = len(components) // 4
+    verify_fn(self, lambda: self._build_dataset(components), num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/take_test.py

@@ -0,0 +1,185 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.take()`."""
+
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class TakeTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(count=[-1, 0, 4, 10, 25])))
+  def testBasic(self, count):
+    components = (np.arange(10),)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).take(count)
+    self.assertEqual(
+        [c.shape[1:] for c in components],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    num_output = min(count, 10) if count != -1 else 10
+    self.assertDatasetProduces(
+        dataset, [tuple(components[0][i:i + 1]) for i in range(num_output)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).take(1, name="take")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class TakeDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  def _build_take_dataset(self, count, options=None):
+    dataset = dataset_ops.Dataset.range(100).take(count)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True]),
+          combinations.combine(count=[50], num_outputs=[50]) +
+          combinations.combine(count=[200, 100, -1], num_outputs=[100]) +
+          combinations.combine(count=[0], num_outputs=[0])))
+  def test(self, verify_fn, symbolic_checkpoint, count, num_outputs):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self._build_take_dataset(count, options),
+              num_outputs)
+
+
+class TakeRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 3, 4])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(10).take(3)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-2, 0, 1])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.from_tensor_slices([]).take(5)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(count=[-1, 0, 4, 10, 25])))
+  def testBasic(self, count):
+    dataset = dataset_ops.Dataset.range(10).take(count)
+    num_output = min(count, 10) if count != -1 else 10
+    for i in range(num_output):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)), i)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=num_output))
+
+
+class TakeGlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              count=[20, 200],
+              repetitions=[1, 2],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def test(
+      self,
+      dataset_range: int,
+      count: int,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = dataset.take(count)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    if repetitions > 1:
+      dataset = dataset.repeat(repetitions)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle_each_iteration)
+
+    expected = list(range(0, min(count, dataset_range))) * repetitions
+    dataset_output = self.getDatasetOutput(
+        dataset, requires_initialization=True)
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+    self.assertLen(dataset_output, self.evaluate(dataset.cardinality()))
+
+
+class TakeGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[10],
+              count=[2, 20],
+              repetitions=[1, 2],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def test(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      count: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      dataset = dataset.take(count)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=min(count, dataset_range) * repetitions,
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/take_while_test.py

@@ -0,0 +1,153 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.take_while()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class TakeWhileTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_elements=[14, 15], window_size=[2]) +
+          combinations.combine(num_elements=[100], window_size=[3])))
+  def testTakeWhileDataset(self, num_elements, window_size):
+
+    def _predicate_func(elem):
+      return array_ops.shape(elem)[0] > (window_size - 1)
+
+    dataset = dataset_ops.Dataset.range(num_elements).batch(window_size)
+    dataset = dataset.take_while(predicate=_predicate_func).flat_map(
+        dataset_ops.Dataset.from_tensor_slices)
+
+    expected_num_elements = int(num_elements / window_size) * window_size
+    self.assertDatasetProduces(dataset, np.arange(expected_num_elements))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_elements=[10], upper_bound=[2]) +
+          combinations.combine(num_elements=[16], upper_bound=[7]) +
+          combinations.combine(num_elements=[100], upper_bound=[99]) +
+          combinations.combine(num_elements=[100], upper_bound=[101]) +
+          combinations.combine(num_elements=[0], upper_bound=[1])))
+  def testTakeWhileDatasetRange(self, num_elements, upper_bound):
+    dataset = dataset_ops.Dataset.range(num_elements).take_while(
+        lambda x: x < upper_bound)
+
+    self.assertDatasetProduces(dataset,
+                               np.arange(min(num_elements, upper_bound)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTakeWhileDatasetString(self):
+
+    def not_equal(string):
+      return lambda x: math_ops.not_equal(x, constant_op.constant(string))
+
+    string = ["this", "is", "the", "test", "for", "strings"]
+    dataset = dataset_ops.Dataset.from_tensor_slices(string).take_while(
+        predicate=not_equal("test"))
+
+    next_element = self.getNext(dataset)
+    self.assertEqual(b"this", self.evaluate(next_element()))
+    self.assertEqual(b"is", self.evaluate(next_element()))
+    self.assertEqual(b"the", self.evaluate(next_element()))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(size=[5], index=[3]) +
+          combinations.combine(size=[10], index=[0]) +
+          combinations.combine(size=[100], index=[5]) +
+          combinations.combine(size=[8], index=[7])))
+  def testTakewhileDatasetShortCircuit(self, size, index):
+
+    def _predicate_func(data_elem):
+      return data_elem
+
+    boolean_array = [True] * size
+    boolean_array[index] = False
+    dataset = dataset_ops.Dataset.from_tensor_slices(boolean_array).take_while(
+        predicate=_predicate_func)
+
+    next_element = self.getNext(dataset)
+
+    for _ in range(index):
+      self.assertTrue(self.evaluate(next_element()))
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTakeWhileDatasetWithRepeat(self):
+    dataset = dataset_ops.Dataset.range(10).take_while(
+        predicate=lambda x: x < 2).repeat(5)
+    self.assertDatasetProduces(dataset, np.tile([0, 1], 5))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTakeWhileDatasetStops(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.take_while(
+        lambda x: math_ops.logical_not(math_ops.equal(x, 5)))
+    self.assertDatasetProduces(dataset, range(5))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).take_while(
+        lambda _: True, name="take_while")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class TakeWhileCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                              parameterized.TestCase):
+
+  def _build_dataset(self, num_elements, upper_bound, options=None):
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.take_while(predicate=lambda x: x < upper_bound)
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True]),
+          combinations.combine(num_elements=[10, 23], upper_bound=[10, 23])))
+  def test(self, verify_fn, symbolic_checkpoint, num_elements, upper_bound):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self,
+              lambda: self._build_dataset(num_elements, upper_bound, options),
+              min(num_elements, upper_bound))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/test_base.py

@@ -0,0 +1,447 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test utilities for tf.data functionality."""
+import os
+import random
+import re
+
+from tensorflow.python.data.experimental.ops import lookup_ops as data_lookup_ops
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import test_mode
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+def default_test_combinations():
+  """Returns the default test combinations for tf.data tests."""
+  return combinations.combine(tf_api_version=[1, 2], mode=["eager", "graph"])
+
+
+def eager_only_combinations():
+  """Returns the default test combinations for eager mode only tf.data tests."""
+  return combinations.combine(tf_api_version=[1, 2], mode="eager")
+
+
+def graph_only_combinations():
+  """Returns the default test combinations for graph mode only tf.data tests."""
+  return combinations.combine(tf_api_version=[1, 2], mode="graph")
+
+
+def v1_only_combinations():
+  """Returns the default test combinations for v1 only tf.data tests."""
+  return combinations.combine(tf_api_version=1, mode=["eager", "graph"])
+
+
+def v2_only_combinations():
+  """Returns the default test combinations for v2 only tf.data tests."""
+  return combinations.combine(tf_api_version=2, mode=["eager", "graph"])
+
+
+def v2_eager_only_combinations():
+  """Returns the default test combinations for v2 eager only tf.data tests."""
+  return combinations.combine(tf_api_version=2, mode="eager")
+
+
+class DatasetTestBase(test.TestCase):
+  """Base class for dataset tests."""
+
+  def setUp(self):
+    super().setUp()
+    test_mode.toggle_test_mode(True)
+
+  def assert_op_cancelled(self, op):
+    with self.assertRaises(errors.CancelledError):
+      self.evaluate(op)
+
+  def assertValuesEqual(self, expected, actual):
+    """Asserts that two values are equal."""
+    if isinstance(expected, dict):
+      self.assertItemsEqual(list(expected.keys()), list(actual.keys()))
+      for k in expected.keys():
+        self.assertValuesEqual(expected[k], actual[k])
+    elif sparse_tensor.is_sparse(expected):
+      self.assertAllEqual(expected.indices, actual.indices)
+      self.assertAllEqual(expected.values, actual.values)
+      self.assertAllEqual(expected.dense_shape, actual.dense_shape)
+    else:
+      self.assertAllEqual(expected, actual)
+
+  def getNext(self, dataset, requires_initialization=False, shared_name=None):
+    """Returns a callable that returns the next element of the dataset.
+
+    Example use:
+    ```python
+    # In both graph and eager modes
+    dataset = ...
+    get_next = self.getNext(dataset)
+    result = self.evaluate(get_next())
+    ```
+
+    Args:
+      dataset: A dataset whose elements will be returned.
+      requires_initialization: Indicates that when the test is executed in graph
+        mode, it should use an initializable iterator to iterate through the
+        dataset (e.g. when it contains stateful nodes). Defaults to False.
+      shared_name: (Optional.) If non-empty, the returned iterator will be
+        shared under the given name across multiple sessions that share the same
+        devices (e.g. when using a remote server).
+    Returns:
+      A callable that returns the next element of `dataset`. Any `TensorArray`
+      objects `dataset` outputs are stacked.
+    """
+    def ta_wrapper(gn):
+      def _wrapper():
+        r = gn()
+        if isinstance(r, tensor_array_ops.TensorArray):
+          return r.stack()
+        else:
+          return r
+      return _wrapper
+
+    # Create an anonymous iterator if we are in eager-mode or are graph inside
+    # of a tf.function.
+    if context.executing_eagerly() or ops.inside_function():
+      iterator = iter(dataset)
+      return ta_wrapper(iterator._next_internal)  # pylint: disable=protected-access
+    else:
+      if requires_initialization:
+        iterator = dataset_ops.make_initializable_iterator(dataset, shared_name)
+        self.evaluate(iterator.initializer)
+      else:
+        iterator = dataset_ops.make_one_shot_iterator(dataset)
+      get_next = iterator.get_next()
+      return ta_wrapper(lambda: get_next)
+
+  def _compareOutputToExpected(self, result_values, expected_values,
+                               assert_items_equal):
+    if assert_items_equal:
+      # TODO(shivaniagrawal): add support for nested elements containing sparse
+      # tensors when needed.
+      self.assertItemsEqual(result_values, expected_values)
+      return
+    for i in range(len(result_values)):
+      nest.assert_same_structure(result_values[i], expected_values[i])
+      for result_value, expected_value in zip(
+          nest.flatten(result_values[i]), nest.flatten(expected_values[i])):
+        self.assertValuesEqual(expected_value, result_value)
+
+  def getDatasetOutput(self, dataset, requires_initialization=False):
+    get_next = self.getNext(
+        dataset, requires_initialization=requires_initialization)
+    return self.getIteratorOutput(get_next)
+
+  def getIteratorOutput(self, get_next):
+    """Evaluates `get_next` until end of input, returning the results."""
+    results = []
+    while True:
+      try:
+        results.append(self.evaluate(get_next()))
+      except errors.OutOfRangeError:
+        break
+    return results
+
+  def assertDatasetProduces(self,
+                            dataset,
+                            expected_output=None,
+                            expected_shapes=None,
+                            expected_error=None,
+                            requires_initialization=False,
+                            num_test_iterations=1,
+                            assert_items_equal=False,
+                            expected_error_iter=1):
+    """Asserts that a dataset produces the expected output / error.
+
+    Args:
+      dataset: A dataset to check for the expected output / error.
+      expected_output: A list of elements that the dataset is expected to
+        produce.
+      expected_shapes: A list of TensorShapes which is expected to match
+        output_shapes of dataset.
+      expected_error: A tuple `(type, predicate)` identifying the expected error
+        `dataset` should raise. The `type` should match the expected exception
+        type, while `predicate` should either be 1) a unary function that inputs
+        the raised exception and returns a boolean indicator of success or 2) a
+        regular expression that is expected to match the error message
+        partially.
+      requires_initialization: Indicates that when the test is executed in graph
+        mode, it should use an initializable iterator to iterate through the
+        dataset (e.g. when it contains stateful nodes). Defaults to False.
+      num_test_iterations: Number of times `dataset` will be iterated. Defaults
+        to 1.
+      assert_items_equal: Tests expected_output has (only) the same elements
+        regardless of order.
+      expected_error_iter: How many times to iterate before expecting an error,
+        if an error is expected.
+    """
+    self.assertTrue(
+        expected_error is not None or expected_output is not None,
+        "Exactly one of expected_output or expected error should be provided.")
+    if expected_error:
+      self.assertTrue(
+          expected_output is None,
+          "Exactly one of expected_output or expected error should be provided."
+      )
+      with self.assertRaisesWithPredicateMatch(expected_error[0],
+                                               expected_error[1]):
+        get_next = self.getNext(
+            dataset, requires_initialization=requires_initialization)
+        for _ in range(expected_error_iter):
+          self.evaluate(get_next())
+      return
+    if expected_shapes:
+      self.assertEqual(expected_shapes,
+                       dataset_ops.get_legacy_output_shapes(dataset))
+    self.assertGreater(num_test_iterations, 0)
+    for _ in range(num_test_iterations):
+      get_next = self.getNext(
+          dataset, requires_initialization=requires_initialization)
+      result = []
+      for _ in range(len(expected_output)):
+        try:
+          result.append(self.evaluate(get_next()))
+        except errors.OutOfRangeError:
+          raise AssertionError(
+              "Dataset ended early, producing %d elements out of %d. "
+              "Dataset output: %s" %
+              (len(result), len(expected_output), str(result)))
+      self._compareOutputToExpected(result, expected_output, assert_items_equal)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+  def assertDatasetsEqual(self, dataset1, dataset2):
+    """Checks that datasets are equal. Supports both graph and eager mode."""
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(dataset1),
+            dataset_ops.get_structure(dataset2)))
+
+    flattened_types = nest.flatten(
+        dataset_ops.get_legacy_output_types(dataset1))
+
+    next1 = self.getNext(dataset1)
+    next2 = self.getNext(dataset2)
+
+    while True:
+      try:
+        op1 = self.evaluate(next1())
+      except errors.OutOfRangeError:
+        with self.assertRaises(errors.OutOfRangeError):
+          self.evaluate(next2())
+        break
+      op2 = self.evaluate(next2())
+
+      op1 = nest.flatten(op1)
+      op2 = nest.flatten(op2)
+      assert len(op1) == len(op2)
+      for i in range(len(op1)):
+        if sparse_tensor.is_sparse(op1[i]) or ragged_tensor.is_ragged(op1[i]):
+          self.assertValuesEqual(op1[i], op2[i])
+        elif flattened_types[i] == dtypes.string:
+          self.assertAllEqual(op1[i], op2[i])
+        else:
+          self.assertAllClose(op1[i], op2[i])
+
+  def assertDatasetsRaiseSameError(self,
+                                   dataset1,
+                                   dataset2,
+                                   exception_class,
+                                   replacements=None):
+    """Checks that datasets raise the same error on the first get_next call."""
+    if replacements is None:
+      replacements = []
+    next1 = self.getNext(dataset1)
+    next2 = self.getNext(dataset2)
+    try:
+      self.evaluate(next1())
+      raise ValueError(
+          "Expected dataset to raise an error of type %s, but it did not." %
+          repr(exception_class))
+    except exception_class as e:
+      expected_message = e.message
+      for old, new, count in replacements:
+        expected_message = expected_message.replace(old, new, count)
+      # Check that the first segment of the error messages are the same.
+      with self.assertRaisesRegex(exception_class, re.escape(expected_message)):
+        self.evaluate(next2())
+
+  def structuredDataset(self, dataset_structure, shape=None,
+                        dtype=dtypes.int64):
+    """Returns a singleton dataset with the given structure."""
+    if shape is None:
+      shape = []
+    if dataset_structure is None:
+      return dataset_ops.Dataset.from_tensors(
+          array_ops.zeros(shape, dtype=dtype))
+    else:
+      return dataset_ops.Dataset.zip(
+          tuple([
+              self.structuredDataset(substructure, shape, dtype)
+              for substructure in dataset_structure
+          ]))
+
+  def verifyRandomAccess(self, dataset, expected):
+    self.verifyRandomAccessInfiniteCardinality(dataset, expected)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=len(expected)))
+
+  def verifyRandomAccessInfiniteCardinality(self, dataset, expected):
+    """Tests randomly accessing elements of a dataset."""
+    # Tests accessing the elements in a shuffled order with repeats.
+    len_expected = len(expected)
+    indices = list(range(len_expected)) * 2
+    random.shuffle(indices)
+    for i in indices:
+      self.assertAllEqual(expected[i],
+                          self.evaluate(random_access.at(dataset, i)))
+
+    # Tests accessing the elements in order.
+    indices = set(sorted(indices))
+    for i in indices:
+      self.assertAllEqual(expected[i],
+                          self.evaluate(random_access.at(dataset, i)))
+
+  def textFileInitializer(self, vals):
+    file = os.path.join(self.get_temp_dir(), "text_file_initializer")
+    with open(file, "w") as f:
+      f.write("\n".join(str(v) for v in vals) + "\n")
+    return lookup_ops.TextFileInitializer(file, dtypes.int64,
+                                          lookup_ops.TextFileIndex.LINE_NUMBER,
+                                          dtypes.int64,
+                                          lookup_ops.TextFileIndex.WHOLE_LINE)
+
+  def keyValueTensorInitializer(self, vals):
+    keys_tensor = constant_op.constant(
+        list(range(len(vals))), dtype=dtypes.int64)
+    vals_tensor = constant_op.constant(vals)
+    return lookup_ops.KeyValueTensorInitializer(keys_tensor, vals_tensor)
+
+  def datasetInitializer(self, vals):
+    keys = dataset_ops.Dataset.range(len(vals))
+    values = dataset_ops.Dataset.from_tensor_slices(vals)
+    ds = dataset_ops.Dataset.zip((keys, values))
+    return data_lookup_ops.DatasetInitializer(ds)
+
+  def lookupTableInitializer(self, init_source, vals):
+    """Returns a lookup table initializer for the given source and values.
+
+    Args:
+      init_source: One of ["textfile", "keyvalue", "dataset"], indicating what
+        type of initializer to use.
+      vals: The initializer values. The keys will be `range(len(vals))`.
+    """
+    if init_source == "textfile":
+      return self.textFileInitializer(vals)
+    elif init_source == "keyvaluetensor":
+      return self.keyValueTensorInitializer(vals)
+    elif init_source == "dataset":
+      return self.datasetInitializer(vals)
+    else:
+      raise ValueError("Unrecognized init_source: " + init_source)
+
+  def graphRoundTrip(self, dataset, allow_stateful=False):
+    """Converts a dataset to a graph and back."""
+    graph = gen_dataset_ops.dataset_to_graph(
+        dataset._variant_tensor, allow_stateful=allow_stateful)  # pylint: disable=protected-access
+    return dataset_ops.from_variant(
+        gen_experimental_dataset_ops.dataset_from_graph(graph),
+        dataset.element_spec)
+
+  def structuredElement(self, element_structure, shape=None,
+                        dtype=dtypes.int64):
+    """Returns an element with the given structure."""
+    if shape is None:
+      shape = []
+    if element_structure is None:
+      return array_ops.zeros(shape, dtype=dtype)
+    else:
+      return tuple([
+          self.structuredElement(substructure, shape, dtype)
+          for substructure in element_structure
+      ])
+
+  def checkDeterminism(self, dataset_fn, expect_determinism, expected_elements):
+    """Tests whether a dataset produces its elements deterministically.
+
+    `dataset_fn` takes a delay_ms argument, which tells it how long to delay
+    production of the first dataset element. This gives us a way to trigger
+    out-of-order production of dataset elements.
+
+    Args:
+      dataset_fn: A function taking a delay_ms argument.
+      expect_determinism: Whether to expect deterministic ordering.
+      expected_elements: The elements expected to be produced by the dataset,
+        assuming the dataset produces elements in deterministic order.
+    """
+    if expect_determinism:
+      dataset = dataset_fn(100)
+      actual = self.getDatasetOutput(dataset)
+      self.assertAllEqual(expected_elements, actual)
+      return
+
+    # We consider the test a success if it succeeds under any delay_ms. The
+    # delay_ms needed to observe non-deterministic ordering varies across
+    # test machines. Usually 10 or 100 milliseconds is enough, but on slow
+    # machines it could take longer.
+    for delay_ms in [10, 100, 1000, 20000, 100000]:
+      dataset = dataset_fn(delay_ms)
+      actual = self.getDatasetOutput(dataset)
+      self.assertCountEqual(expected_elements, actual)
+      for i in range(len(actual)):
+        if actual[i] != expected_elements[i]:
+          return
+    self.fail("Failed to observe nondeterministic ordering")
+
+  def configureDevicesForMultiDeviceTest(self, num_devices):
+    """Configures number of logical devices for multi-device tests.
+
+    It returns a list of device names. If invoked in GPU-enabled runtime, the
+    last device name will be for a GPU device. Otherwise, all device names will
+    be for a CPU device.
+
+    Args:
+      num_devices: The number of devices to configure.
+
+    Returns:
+      A list of device names to use for a multi-device test.
+    """
+    cpus = config.list_physical_devices("CPU")
+    gpus = config.list_physical_devices("GPU")
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration() for _ in range(num_devices)
+    ])
+    devices = ["/device:CPU:" + str(i) for i in range(num_devices - 1)]
+    if gpus:
+      devices.append("/device:GPU:0")
+    else:
+      devices.append("/device:CPU:" + str(num_devices - 1))
+    return devices

tensorflow/python/data/kernel_tests/text_line_dataset_test.py

@@ -0,0 +1,274 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.TextLineDataset`."""
+import gzip
+import os
+import pathlib
+import zlib
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+try:
+  import psutil  # pylint: disable=g-import-not-at-top
+  psutil_import_succeeded = True
+except ImportError:
+  psutil_import_succeeded = False
+
+
+class TextLineDatasetTestBase(test_base.DatasetTestBase):
+  """Base class for setting up and testing TextLineDataset."""
+
+  def _lineText(self, f, l):
+    return compat.as_bytes("%d: %d" % (f, l))
+
+  def _createFiles(self,
+                   num_files,
+                   num_lines,
+                   crlf=False,
+                   compression_type=None):
+    filenames = []
+    for i in range(num_files):
+      fn = os.path.join(self.get_temp_dir(), "text_line.%d.txt" % i)
+      filenames.append(fn)
+      contents = []
+      for j in range(num_lines):
+        contents.append(self._lineText(i, j))
+        # Always include a newline after the record unless it is
+        # at the end of the file, in which case we include it
+        if j + 1 != num_lines or i == 0:
+          contents.append(b"\r\n" if crlf else b"\n")
+      contents = b"".join(contents)
+
+      if not compression_type:
+        with open(fn, "wb") as f:
+          f.write(contents)
+      elif compression_type == "GZIP":
+        with gzip.GzipFile(fn, "wb") as f:
+          f.write(contents)
+      elif compression_type == "ZLIB":
+        contents = zlib.compress(contents)
+        with open(fn, "wb") as f:
+          f.write(contents)
+      else:
+        raise ValueError("Unsupported compression_type", compression_type)
+
+    return filenames
+
+
+class TextLineDatasetTest(TextLineDatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(compression_type=[None, "GZIP", "ZLIB"])))
+  def testBasic(self, compression_type):
+    test_filenames = self._createFiles(
+        2, 5, crlf=True, compression_type=compression_type)
+
+    def dataset_fn(filenames, num_epochs, batch_size=None):
+      repeat_dataset = readers.TextLineDataset(
+          filenames, compression_type=compression_type).repeat(num_epochs)
+      if batch_size:
+        return repeat_dataset.batch(batch_size)
+      return repeat_dataset
+
+    # Basic test: read from file 0.
+    expected_output = [self._lineText(0, i) for i in range(5)]
+    self.assertDatasetProduces(
+        dataset_fn([test_filenames[0]], 1), expected_output=expected_output)
+
+    # Basic test: read from file 1.
+    self.assertDatasetProduces(
+        dataset_fn([test_filenames[1]], 1),
+        expected_output=[self._lineText(1, i) for i in range(5)])
+
+    # Basic test: read from both files.
+    expected_output = [self._lineText(0, i) for i in range(5)]
+    expected_output.extend(self._lineText(1, i) for i in range(5))
+    self.assertDatasetProduces(
+        dataset_fn(test_filenames, 1), expected_output=expected_output)
+
+    # Test repeated iteration through both files.
+    expected_output = [self._lineText(0, i) for i in range(5)]
+    expected_output.extend(self._lineText(1, i) for i in range(5))
+    self.assertDatasetProduces(
+        dataset_fn(test_filenames, 10), expected_output=expected_output * 10)
+
+    # Test batched and repeated iteration through both files.
+    self.assertDatasetProduces(
+        dataset_fn(test_filenames, 10, 5),
+        expected_output=[[self._lineText(0, i) for i in range(5)],
+                         [self._lineText(1, i) for i in range(5)]] * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelRead(self):
+    test_filenames = self._createFiles(10, 10)
+    files = dataset_ops.Dataset.from_tensor_slices(test_filenames).repeat(10)
+    expected_output = []
+    for j in range(10):
+      expected_output.extend(self._lineText(j, i) for i in range(10))
+    dataset = readers.TextLineDataset(files, num_parallel_reads=4)
+    self.assertDatasetProduces(
+        dataset, expected_output=expected_output * 10, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBuffering(self):
+    test_filenames = self._createFiles(2, 5, crlf=True)
+
+    repeat_dataset = readers.TextLineDataset(test_filenames, buffer_size=10)
+    expected_output = []
+    for j in range(2):
+      expected_output.extend([self._lineText(j, i) for i in range(5)])
+    self.assertDatasetProduces(repeat_dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testIteratorResourceCleanup(self):
+    filename = os.path.join(self.get_temp_dir(), "text.txt")
+    with open(filename, "wt") as f:
+      for i in range(3):
+        f.write("%d\n" % (i,))
+    first_iterator = iter(readers.TextLineDataset(filename))
+    self.assertEqual(b"0", next(first_iterator).numpy())
+    second_iterator = iter(readers.TextLineDataset(filename))
+    self.assertEqual(b"0", next(second_iterator).numpy())
+    # Eager kernel caching is based on op attributes, which includes the
+    # Dataset's output shape. Create a different kernel to test that they
+    # don't create resources with the same names.
+    different_kernel_iterator = iter(
+        readers.TextLineDataset(filename).repeat().batch(16))
+    self.assertEqual([16], next(different_kernel_iterator).shape)
+    # Remove our references to the Python Iterator objects, which (assuming no
+    # reference cycles) is enough to trigger DestroyResourceOp and close the
+    # partially-read files.
+    del first_iterator
+    del second_iterator
+    del different_kernel_iterator
+    if not psutil_import_succeeded:
+      self.skipTest(
+          "psutil is required to check that we've closed our files.")
+    open_files = psutil.Process().open_files()
+    self.assertNotIn(filename, [open_file.path for open_file in open_files])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPathlib(self):
+    files = self._createFiles(1, 5)
+    files = [pathlib.Path(f) for f in files]
+
+    expected_output = [self._lineText(0, i) for i in range(5)]
+    ds = readers.TextLineDataset(files)
+    self.assertDatasetProduces(
+        ds, expected_output=expected_output, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    files = self._createFiles(1, 5)
+    expected_output = [self._lineText(0, i) for i in range(5)]
+    ds = readers.TextLineDataset(files, name="text_line_dataset")
+    self.assertDatasetProduces(
+        ds, expected_output=expected_output, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyFileList(self):
+    dataset = readers.TextLineDataset(filenames=[])
+    self.assertDatasetProduces(dataset, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileDoesNotExist(self):
+    dataset = readers.TextLineDataset(filenames=["File not exist"])
+    with self.assertRaisesRegex(errors.NotFoundError,
+                                "No such file or directory"):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileNamesMustBeStrings(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "The `filenames` argument must contain `tf.string` elements. Got "
+        "`tf.int32` elements."):
+      readers.TextLineDataset(filenames=0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileNamesDatasetMustContainStrings(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "The `filenames` argument must contain `tf.string` elements. Got a "
+        "dataset of `tf.int32` elements."):
+      filenames = dataset_ops.Dataset.from_tensors(0)
+      readers.TextLineDataset(filenames)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileNamesMustBeScalars(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "The `filenames` argument must contain `tf.string` elements of shape "
+        r"\[\] \(i.e. scalars\)."):
+      filenames = dataset_ops.Dataset.from_tensors([["File 1", "File 2"],
+                                                    ["File 3", "File 4"]])
+      readers.TextLineDataset(filenames)
+
+
+class TextLineDatasetCheckpointTest(TextLineDatasetTestBase,
+                                    checkpoint_test_base.CheckpointTestBase,
+                                    parameterized.TestCase):
+
+  def _build_iterator_graph(
+      self, test_filenames, symbolic_checkpoint, compression_type=None
+  ):
+    dataset = readers.TextLineDataset(
+        test_filenames, compression_type=compression_type, buffer_size=10
+    )
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              symbolic_checkpoint=[False, True],
+              compression_type=[None, "GZIP", "ZLIB"],
+          ),
+      )
+  )
+  def test(self, verify_fn, symbolic_checkpoint, compression_type):
+    num_files = 5
+    lines_per_file = 5
+    num_outputs = num_files * lines_per_file
+    test_filenames = self._createFiles(
+        num_files, lines_per_file, crlf=True, compression_type=compression_type)
+    verify_fn(
+        self,
+        lambda: self._build_iterator_graph(
+            test_filenames, symbolic_checkpoint, compression_type
+        ),
+        num_outputs,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/tf_record_dataset_test.py

@@ -0,0 +1,275 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.TFRecordDataset`."""
+import gzip
+import os
+import pathlib
+import zlib
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.kernel_tests import tf_record_test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
+
+
+class TFRecordDatasetTest(tf_record_test_base.TFRecordTestBase,
+                          parameterized.TestCase):
+
+  def _dataset_factory(self,
+                       filenames,
+                       compression_type="",
+                       num_epochs=1,
+                       batch_size=None):
+
+    repeat_dataset = readers.TFRecordDataset(
+        filenames, compression_type).repeat(num_epochs)
+    if batch_size:
+      return repeat_dataset.batch(batch_size)
+    return repeat_dataset
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConstructorErrorsTensorInput(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "The `filenames` argument must contain `tf.string` elements. Got "
+        "`tf.int32` elements."):
+      readers.TFRecordDataset([1, 2, 3])
+    with self.assertRaisesRegex(
+        TypeError,
+        "The `filenames` argument must contain `tf.string` elements. Got "
+        "`tf.int32` elements."):
+      readers.TFRecordDataset(constant_op.constant([1, 2, 3]))
+    # convert_to_tensor raises different errors in graph and eager
+    with self.assertRaises(Exception):
+      readers.TFRecordDataset(object())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadOneEpoch(self):
+    # Basic test: read from file 0.
+    dataset = self._dataset_factory(self._filenames[0])
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[self._record(0, i) for i in range(self._num_records)])
+
+    # Basic test: read from file 1.
+    dataset = self._dataset_factory(self._filenames[1])
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[self._record(1, i) for i in range(self._num_records)])
+
+    # Basic test: read from both files.
+    dataset = self._dataset_factory(self._filenames)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadTenEpochs(self):
+    dataset = self._dataset_factory(self._filenames, num_epochs=10)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadTenEpochsOfBatches(self):
+    dataset = self._dataset_factory(
+        self._filenames, num_epochs=10, batch_size=self._num_records)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.append(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadZlibFiles(self):
+    zlib_files = []
+    for i, fn in enumerate(self._filenames):
+      with open(fn, "rb") as f:
+        cdata = zlib.compress(f.read())
+
+        zfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.z" % i)
+        with open(zfn, "wb") as f:
+          f.write(cdata)
+        zlib_files.append(zfn)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    dataset = self._dataset_factory(zlib_files, compression_type="ZLIB")
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadGzipFiles(self):
+    gzip_files = []
+    for i, fn in enumerate(self._filenames):
+      with open(fn, "rb") as f:
+        gzfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.gz" % i)
+        with gzip.GzipFile(gzfn, "wb") as gzf:
+          gzf.write(f.read())
+        gzip_files.append(gzfn)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    dataset = self._dataset_factory(gzip_files, compression_type="GZIP")
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadWithBuffer(self):
+    one_mebibyte = 2**20
+    dataset = readers.TFRecordDataset(
+        self._filenames, buffer_size=one_mebibyte)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadFromDatasetOfFiles(self):
+    files = dataset_ops.Dataset.from_tensor_slices(self._filenames)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    dataset = readers.TFRecordDataset(files)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadTenEpochsFromDatasetOfFilesInParallel(self):
+    files = dataset_ops.Dataset.from_tensor_slices(
+        self._filenames).repeat(10)
+    expected_output = []
+    for j in range(self._num_files):
+      expected_output.extend(
+          [self._record(j, i) for i in range(self._num_records)])
+    dataset = readers.TFRecordDataset(files, num_parallel_reads=4)
+    self.assertDatasetProduces(
+        dataset, expected_output=expected_output * 10, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPathlib(self):
+    files = [pathlib.Path(self._filenames[0])]
+
+    expected_output = [self._record(0, i) for i in range(self._num_records)]
+    ds = readers.TFRecordDataset(files)
+    self.assertDatasetProduces(
+        ds, expected_output=expected_output, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    files = [self._filenames[0]]
+
+    expected_output = [self._record(0, i) for i in range(self._num_records)]
+    ds = readers.TFRecordDataset(files, name="tf_record_dataset")
+    self.assertDatasetProduces(
+        ds, expected_output=expected_output, assert_items_equal=True)
+
+
+class TFRecordDatasetCheckpointTest(tf_record_test_base.TFRecordTestBase,
+                                    checkpoint_test_base.CheckpointTestBase,
+                                    parameterized.TestCase):
+
+  def make_dataset(self,
+                   num_epochs,
+                   compression_type=None,
+                   buffer_size=None,
+                   symbolic_checkpoint=False):
+    filenames = self._createFiles()
+    if compression_type == "ZLIB":
+      zlib_files = []
+      for i, fn in enumerate(filenames):
+        with open(fn, "rb") as f:
+          cdata = zlib.compress(f.read())
+          zfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.z" % i)
+          with open(zfn, "wb") as f:
+            f.write(cdata)
+          zlib_files.append(zfn)
+      filenames = zlib_files
+
+    elif compression_type == "GZIP":
+      gzip_files = []
+      for i, fn in enumerate(self._filenames):
+        with open(fn, "rb") as f:
+          gzfn = os.path.join(self.get_temp_dir(), "tfrecord_%s.gz" % i)
+          with gzip.GzipFile(gzfn, "wb") as gzf:
+            gzf.write(f.read())
+          gzip_files.append(gzfn)
+      filenames = gzip_files
+
+    dataset = readers.TFRecordDataset(
+        filenames, compression_type,
+        buffer_size=buffer_size).repeat(num_epochs)
+    if symbolic_checkpoint:
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[True, False])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    num_epochs = 5
+    num_outputs = num_epochs * self._num_files * self._num_records
+    verify_fn(
+        self,
+        lambda: self.make_dataset(
+            num_epochs, symbolic_checkpoint=symbolic_checkpoint
+        ),
+        num_outputs,
+    )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(buffer_size=[0, 5])))
+  def testBufferSize(self, verify_fn, buffer_size):
+    num_epochs = 5
+    num_outputs = num_epochs * self._num_files * self._num_records
+    verify_fn(self,
+              lambda: self.make_dataset(num_epochs, buffer_size=buffer_size),
+              num_outputs)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(compression_type=[None, "GZIP", "ZLIB"])))
+  def testCompressionTypes(self, verify_fn, compression_type):
+    num_epochs = 5
+    num_outputs = num_epochs * self._num_files * self._num_records
+    # pylint: disable=g-long-lambda
+    verify_fn(
+        self, lambda: self.make_dataset(
+            num_epochs, compression_type=compression_type), num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/tf_record_test_base.py

@@ -0,0 +1,338 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Base class for testing reader datasets."""
+
+import os
+
+from tensorflow.core.example import example_pb2
+from tensorflow.core.example import feature_pb2
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import readers as core_readers
+from tensorflow.python.framework import dtypes
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.util import compat
+
+
+class FeaturesTestBase(test_base.DatasetTestBase):
+  """Base class for testing TFRecord-based features."""
+
+  def setUp(self):
+    super(FeaturesTestBase, self).setUp()
+    self._num_files = 2
+    self._num_records = 7
+    self._filenames = self._createFiles()
+
+  def make_batch_feature(self,
+                         filenames,
+                         num_epochs,
+                         batch_size,
+                         label_key=None,
+                         reader_num_threads=1,
+                         parser_num_threads=1,
+                         shuffle=False,
+                         shuffle_seed=None,
+                         drop_final_batch=False):
+    self.filenames = filenames
+    self.num_epochs = num_epochs
+    self.batch_size = batch_size
+
+    return readers.make_batched_features_dataset(
+        file_pattern=self.filenames,
+        batch_size=self.batch_size,
+        features={
+            "file": parsing_ops.FixedLenFeature([], dtypes.int64),
+            "record": parsing_ops.FixedLenFeature([], dtypes.int64),
+            "keywords": parsing_ops.VarLenFeature(dtypes.string),
+            "label": parsing_ops.FixedLenFeature([], dtypes.string),
+        },
+        label_key=label_key,
+        reader=core_readers.TFRecordDataset,
+        num_epochs=self.num_epochs,
+        shuffle=shuffle,
+        shuffle_seed=shuffle_seed,
+        reader_num_threads=reader_num_threads,
+        parser_num_threads=parser_num_threads,
+        drop_final_batch=drop_final_batch)
+
+  def _record(self, f, r, l):
+    example = example_pb2.Example(
+        features=feature_pb2.Features(
+            feature={
+                "file":
+                    feature_pb2.Feature(
+                        int64_list=feature_pb2.Int64List(value=[f])),
+                "record":
+                    feature_pb2.Feature(
+                        int64_list=feature_pb2.Int64List(value=[r])),
+                "keywords":
+                    feature_pb2.Feature(
+                        bytes_list=feature_pb2.BytesList(
+                            value=self._get_keywords(f, r))),
+                "label":
+                    feature_pb2.Feature(
+                        bytes_list=feature_pb2.BytesList(
+                            value=[compat.as_bytes(l)]))
+            }))
+    return example.SerializeToString()
+
+  def _get_keywords(self, f, r):
+    num_keywords = 1 + (f + r) % 2
+    keywords = []
+    for index in range(num_keywords):
+      keywords.append(compat.as_bytes("keyword%d" % index))
+    return keywords
+
+  def _sum_keywords(self, num_files):
+    sum_keywords = 0
+    for i in range(num_files):
+      for j in range(self._num_records):
+        sum_keywords += 1 + (i + j) % 2
+    return sum_keywords
+
+  def _createFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i)
+      filenames.append(fn)
+      writer = python_io.TFRecordWriter(fn)
+      for j in range(self._num_records):
+        writer.write(self._record(i, j, "fake-label"))
+      writer.close()
+    return filenames
+
+  def _run_actual_batch(self, outputs, label_key_provided=False):
+    if label_key_provided:
+      # outputs would be a tuple of (feature dict, label)
+      features, label = self.evaluate(outputs())
+    else:
+      features = self.evaluate(outputs())
+      label = features["label"]
+    file_out = features["file"]
+    keywords_indices = features["keywords"].indices
+    keywords_values = features["keywords"].values
+    keywords_dense_shape = features["keywords"].dense_shape
+    record = features["record"]
+    return ([
+        file_out, keywords_indices, keywords_values, keywords_dense_shape,
+        record, label
+    ])
+
+  def _next_actual_batch(self, label_key_provided=False):
+    return self._run_actual_batch(self.outputs, label_key_provided)
+
+  def _interleave(self, iterators, cycle_length):
+    pending_iterators = iterators
+    open_iterators = []
+    num_open = 0
+    for i in range(cycle_length):
+      if pending_iterators:
+        open_iterators.append(pending_iterators.pop(0))
+        num_open += 1
+
+    while num_open:
+      for i in range(min(cycle_length, len(open_iterators))):
+        if open_iterators[i] is None:
+          continue
+        try:
+          yield next(open_iterators[i])
+        except StopIteration:
+          if pending_iterators:
+            open_iterators[i] = pending_iterators.pop(0)
+          else:
+            open_iterators[i] = None
+            num_open -= 1
+
+  def _next_expected_batch(self,
+                           file_indices,
+                           batch_size,
+                           num_epochs,
+                           cycle_length=1):
+
+    def _next_record(file_indices):
+      for j in file_indices:
+        for i in range(self._num_records):
+          yield j, i, compat.as_bytes("fake-label")
+
+    def _next_record_interleaved(file_indices, cycle_length):
+      return self._interleave([_next_record([i]) for i in file_indices],
+                              cycle_length)
+
+    file_batch = []
+    keywords_batch_indices = []
+    keywords_batch_values = []
+    keywords_batch_max_len = 0
+    record_batch = []
+    batch_index = 0
+    label_batch = []
+    for _ in range(num_epochs):
+      if cycle_length == 1:
+        next_records = _next_record(file_indices)
+      else:
+        next_records = _next_record_interleaved(file_indices, cycle_length)
+      for record in next_records:
+        f = record[0]
+        r = record[1]
+        label_batch.append(record[2])
+        file_batch.append(f)
+        record_batch.append(r)
+        keywords = self._get_keywords(f, r)
+        keywords_batch_values.extend(keywords)
+        keywords_batch_indices.extend(
+            [[batch_index, i] for i in range(len(keywords))])
+        batch_index += 1
+        keywords_batch_max_len = max(keywords_batch_max_len, len(keywords))
+        if len(file_batch) == batch_size:
+          yield [
+              file_batch, keywords_batch_indices, keywords_batch_values,
+              [batch_size, keywords_batch_max_len], record_batch, label_batch
+          ]
+          file_batch = []
+          keywords_batch_indices = []
+          keywords_batch_values = []
+          keywords_batch_max_len = 0
+          record_batch = []
+          batch_index = 0
+          label_batch = []
+    if file_batch:
+      yield [
+          file_batch, keywords_batch_indices, keywords_batch_values,
+          [len(file_batch), keywords_batch_max_len], record_batch, label_batch
+      ]
+
+  def _verify_records(self,
+                      batch_size,
+                      file_index=None,
+                      num_epochs=1,
+                      label_key_provided=False,
+                      interleave_cycle_length=1):
+    if file_index is not None:
+      file_indices = [file_index]
+    else:
+      file_indices = range(self._num_files)
+
+    for expected_batch in self._next_expected_batch(
+        file_indices,
+        batch_size,
+        num_epochs,
+        cycle_length=interleave_cycle_length):
+      actual_batch = self._next_actual_batch(
+          label_key_provided=label_key_provided)
+      for i in range(len(expected_batch)):
+        self.assertAllEqual(expected_batch[i], actual_batch[i])
+
+
+class TFRecordTestBase(test_base.DatasetTestBase):
+  """Base class for TFRecord-based tests."""
+
+  def setUp(self):
+    super(TFRecordTestBase, self).setUp()
+    self._num_files = 2
+    self._num_records = 7
+    self._filenames = self._createFiles()
+
+  def _interleave(self, iterators, cycle_length):
+    pending_iterators = iterators
+    open_iterators = []
+    num_open = 0
+    for i in range(cycle_length):
+      if pending_iterators:
+        open_iterators.append(pending_iterators.pop(0))
+        num_open += 1
+
+    while num_open:
+      for i in range(min(cycle_length, len(open_iterators))):
+        if open_iterators[i] is None:
+          continue
+        try:
+          yield next(open_iterators[i])
+        except StopIteration:
+          if pending_iterators:
+            open_iterators[i] = pending_iterators.pop(0)
+          else:
+            open_iterators[i] = None
+            num_open -= 1
+
+  def _next_expected_batch(self, file_indices, batch_size, num_epochs,
+                           cycle_length, drop_final_batch, use_parser_fn):
+
+    def _next_record(file_indices):
+      for j in file_indices:
+        for i in range(self._num_records):
+          yield j, i
+
+    def _next_record_interleaved(file_indices, cycle_length):
+      return self._interleave([_next_record([i]) for i in file_indices],
+                              cycle_length)
+
+    record_batch = []
+    batch_index = 0
+    for _ in range(num_epochs):
+      if cycle_length == 1:
+        next_records = _next_record(file_indices)
+      else:
+        next_records = _next_record_interleaved(file_indices, cycle_length)
+      for f, r in next_records:
+        record = self._record(f, r)
+        if use_parser_fn:
+          record = record[1:]
+        record_batch.append(record)
+        batch_index += 1
+        if len(record_batch) == batch_size:
+          yield record_batch
+          record_batch = []
+          batch_index = 0
+    if record_batch and not drop_final_batch:
+      yield record_batch
+
+  def _verify_records(self, outputs, batch_size, file_index, num_epochs,
+                      interleave_cycle_length, drop_final_batch, use_parser_fn):
+    if file_index is not None:
+      if isinstance(file_index, list):
+        file_indices = file_index
+      else:
+        file_indices = [file_index]
+    else:
+      file_indices = range(self._num_files)
+
+    for expected_batch in self._next_expected_batch(
+        file_indices, batch_size, num_epochs, interleave_cycle_length,
+        drop_final_batch, use_parser_fn):
+      actual_batch = self.evaluate(outputs())
+      self.assertAllEqual(expected_batch, actual_batch)
+
+  def _record(self, f, r):
+    return compat.as_bytes("Record %d of file %d" % (r, f))
+
+  def _createFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i)
+      filenames.append(fn)
+      writer = python_io.TFRecordWriter(fn)
+      for j in range(self._num_records):
+        writer.write(self._record(i, j))
+      writer.close()
+    return filenames
+
+  def _writeFile(self, name, data):
+    filename = os.path.join(self.get_temp_dir(), name)
+    writer = python_io.TFRecordWriter(filename)
+    for d in data:
+      writer.write(compat.as_bytes(str(d)))
+    writer.close()
+    return filename
+

tensorflow/python/data/kernel_tests/unbatch_test.py

@@ -0,0 +1,296 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.unbatch()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchWithUnknownRankInput(self):
+    dataset = dataset_ops.Dataset.from_tensors([0, 1, 2, 3]).unbatch()
+    self.assertDatasetProduces(dataset, range(4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchScalarDataset(self):
+    data = tuple([math_ops.range(10) for _ in range(3)])
+    data = dataset_ops.Dataset.from_tensor_slices(data)
+    expected_types = (dtypes.int32,) * 3
+    data = data.batch(2)
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+    data = data.unbatch()
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+
+    self.assertDatasetProduces(data, [(i,) * 3 for i in range(10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchNestedDataset(self):
+    data = dataset_ops.Dataset.from_tensors(
+        [dataset_ops.Dataset.range(10) for _ in range(10)])
+    data = data.unbatch().flat_map(lambda x: x)
+    self.assertDatasetProduces(data, list(range(10)) * 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchDatasetWithStrings(self):
+    data = tuple([math_ops.range(10) for _ in range(3)])
+    data = dataset_ops.Dataset.from_tensor_slices(data)
+    data = data.map(lambda x, y, z: (x, string_ops.as_string(y), z))
+    expected_types = (dtypes.int32, dtypes.string, dtypes.int32)
+    data = data.batch(2)
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+    data = data.unbatch()
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+
+    self.assertDatasetProduces(
+        data, [(i, compat.as_bytes(str(i)), i) for i in range(10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchDatasetWithSparseTensor(self):
+    st = sparse_tensor.SparseTensorValue(
+        indices=[[i, i] for i in range(10)],
+        values=list(range(10)),
+        dense_shape=[10, 10])
+    data = dataset_ops.Dataset.from_tensors(st)
+    data = data.unbatch()
+    data = data.batch(5)
+    data = data.unbatch()
+    expected_output = [
+        sparse_tensor.SparseTensorValue([[i]], [i], [10]) for i in range(10)
+    ]
+    self.assertDatasetProduces(data, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchDatasetWithDenseSparseAndRaggedTensor(self):
+    st = sparse_tensor.SparseTensorValue(
+        indices=[[i, i] for i in range(10)],
+        values=list(range(10)),
+        dense_shape=[10, 10])
+    rt = ragged_factory_ops.constant_value([[[0]], [[1]], [[2]], [[3]], [[4]],
+                                            [[5]], [[6]], [[7]], [[8]], [[9]]])
+    data = dataset_ops.Dataset.from_tensors((list(range(10)), st, rt))
+    data = data.unbatch()
+    data = data.batch(5)
+    data = data.unbatch()
+    expected_output = [(i, sparse_tensor.SparseTensorValue([[i]], [i], [10]),
+                        ragged_factory_ops.constant_value([[i]]))
+                       for i in range(10)]
+    self.assertDatasetProduces(
+        data, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchDatasetWithRaggedTensor(self):
+    rt = ragged_factory_ops.constant_value([[[0]], [[1]], [[2]], [[3]], [[4]],
+                                            [[5]], [[6]], [[7]], [[8]], [[9]]])
+    data = dataset_ops.Dataset.from_tensors(rt)
+    data = data.unbatch()
+    data = data.batch(5)
+    data = data.batch(2)
+    data = data.unbatch()
+    expected_output = [
+        ragged_factory_ops.constant_value([[[0]], [[1]], [[2]], [[3]], [[4]]]),
+        ragged_factory_ops.constant_value([[[5]], [[6]], [[7]], [[8]], [[9]]]),
+    ]
+    self.assertDatasetProduces(
+        data, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchSingleElementTupleDataset(self):
+    data = tuple([(math_ops.range(10),) for _ in range(3)])
+    data = dataset_ops.Dataset.from_tensor_slices(data)
+    expected_types = ((dtypes.int32,),) * 3
+    data = data.batch(2)
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+    data = data.unbatch()
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+
+    self.assertDatasetProduces(data, [((i,),) * 3 for i in range(10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchMultiElementTupleDataset(self):
+    data = tuple([(math_ops.range(10 * i, 10 * i + 10),
+                   array_ops.fill([10], "hi")) for i in range(3)])
+    data = dataset_ops.Dataset.from_tensor_slices(data)
+    expected_types = ((dtypes.int32, dtypes.string),) * 3
+    data = data.batch(2)
+    self.assertAllEqual(expected_types,
+                        dataset_ops.get_legacy_output_types(data))
+    data = data.unbatch()
+    self.assertAllEqual(expected_types,
+                        dataset_ops.get_legacy_output_types(data))
+
+    self.assertDatasetProduces(
+        data,
+        [((i, b"hi"), (10 + i, b"hi"), (20 + i, b"hi")) for i in range(10)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchEmpty(self):
+    data = dataset_ops.Dataset.from_tensors(
+        (constant_op.constant([]), constant_op.constant([], shape=[0, 4]),
+         constant_op.constant([], shape=[0, 4, 0])))
+    data = data.unbatch()
+    self.assertDatasetProduces(data, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchStaticShapeMismatch(self):
+    data = dataset_ops.Dataset.from_tensors((np.arange(7), np.arange(8),
+                                             np.arange(9)))
+    with self.assertRaises(ValueError):
+      data.unbatch()
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testUnbatchDynamicShapeMismatch(self):
+    ph1 = array_ops.placeholder(dtypes.int32, shape=[None])
+    ph2 = array_ops.placeholder(dtypes.int32, shape=None)
+    data = dataset_ops.Dataset.from_tensors((ph1, ph2))
+    data = data.unbatch()
+    iterator = dataset_ops.make_initializable_iterator(data)
+    next_element = iterator.get_next()
+
+    with self.cached_session() as sess:
+      # Mismatch in the 0th dimension.
+      sess.run(
+          iterator.initializer,
+          feed_dict={
+              ph1: np.arange(7).astype(np.int32),
+              ph2: np.arange(8).astype(np.int32)
+          })
+      with self.assertRaises(errors.InvalidArgumentError):
+        self.evaluate(next_element)
+
+      # No 0th dimension (i.e. scalar value) for one component.
+      sess.run(
+          iterator.initializer,
+          feed_dict={
+              ph1: np.arange(7).astype(np.int32),
+              ph2: 7
+          })
+      with self.assertRaises(errors.InvalidArgumentError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnbatchDatasetWithUintDtypes(self):
+    components = (
+        np.tile(np.array([[0], [1], [2], [3]], dtype=np.uint8), 2),
+        np.tile(np.array([[1], [2], [3], [256]], dtype=np.uint16), 2),
+        np.tile(np.array([[2], [3], [4], [65536]], dtype=np.uint32), 2),
+        np.tile(np.array([[3], [4], [5], [4294967296]], dtype=np.uint64), 2),
+    )
+    expected_types = (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64)
+    expected_output = [tuple([c[i] for c in components]) for i in range(4)]
+
+    data = dataset_ops.Dataset.from_tensor_slices(components)
+    data = data.batch(2)
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+
+    data = data.unbatch()
+    self.assertEqual(expected_types, dataset_ops.get_legacy_output_types(data))
+    self.assertDatasetProduces(data, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoneComponent(self):
+    dataset = dataset_ops.Dataset.from_tensors(
+        (list(range(10)), None)).unbatch().map(lambda x, y: x)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors([42]).unbatch(name="unbatch")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class UnbatchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                            parameterized.TestCase):
+
+  def build_dataset(self,
+                    multiplier=15.0,
+                    tensor_slice_len=2,
+                    batch_size=2,
+                    options=None):
+    components = (np.arange(tensor_slice_len), np.array([[1, 2, 3]]) *
+                  np.arange(tensor_slice_len)[:, np.newaxis],
+                  np.array(multiplier) * np.arange(tensor_slice_len))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).batch(
+        batch_size).unbatch()
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint):
+    tensor_slice_len = 8
+    batch_size = 2
+    num_outputs = tensor_slice_len
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self.build_dataset(15.0, tensor_slice_len, batch_size, options),
+        num_outputs)
+
+
+class UnbatchRandomAccessTest(test_base.DatasetTestBase,
+                              parameterized.TestCase):
+  @combinations.generate(test_base.default_test_combinations())
+  def test(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=True)
+    dataset = dataset.unbatch()
+    for i in range(8):
+      self.assertEqual(self.evaluate(random_access.at(dataset, i)), i)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNotDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=False)
+    dataset = dataset.unbatch()
+    with self.assertRaises(errors.FailedPreconditionError):
+      self.evaluate(random_access.at(dataset, 0))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 100])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.batch(4, drop_remainder=True)
+    dataset = dataset.unbatch()
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/unique_test.py

@@ -0,0 +1,108 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.unique()`."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class UniqueTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _testSimpleHelper(self, dtype, test_cases):
+    """Test the `unique()` transformation on a list of test cases.
+
+    Args:
+      dtype: The `dtype` of the elements in each test case.
+      test_cases: A list of pairs of lists. The first component is the test
+        input that will be passed to the transformation; the second component is
+        the expected sequence of outputs from the transformation.
+    """
+
+    # The `current_test_case` will be updated when we loop over `test_cases`
+    # below; declare it here so that the generator can capture it once.
+    current_test_case = []
+    dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,
+                                                 dtype).unique()
+
+    for test_case, expected in test_cases:
+      current_test_case = test_case
+      self.assertDatasetProduces(dataset, [
+          compat.as_bytes(element) if dtype == dtypes.string else element
+          for element in expected
+      ])
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSimpleInt(self):
+    for dtype in [dtypes.int32, dtypes.int64]:
+      self._testSimpleHelper(dtype, [
+          ([], []),
+          ([1], [1]),
+          ([1, 1, 1, 1, 1, 1, 1], [1]),
+          ([1, 1, 1, 1, 0], [1, 0]),
+          ([1, 2, 3, 4], [1, 2, 3, 4]),
+          ([1, 2, 4, 3, 2, 1, 2, 3, 4], [1, 2, 4, 3]),
+          ([[1], [1, 1], [1, 1, 1]], [[1], [1, 1], [1, 1, 1]]),
+          ([[1, 1], [1, 1], [2, 2], [3, 3], [1, 1]], [[1, 1], [2, 2], [3, 3]]),
+      ])
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSimpleString(self):
+    self._testSimpleHelper(dtypes.string, [
+        ([], []),
+        (["hello"], ["hello"]),
+        (["hello", "hello", "hello"], ["hello"]),
+        (["hello", "world"], ["hello", "world"]),
+        (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),
+    ])
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testUnsupportedTypes(self):
+    for dtype in [
+        dtypes.bool, dtypes.double, dtypes.complex64, dtypes.float32,
+        dtypes.float64, dtypes.qint16, dtypes.qint32
+    ]:
+      with self.assertRaises(TypeError):
+        _ = dataset_ops.Dataset.from_generator(lambda: [], dtype).unique()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).unique(name="unique")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class UniqueCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                           parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+
+    def build_dataset(num_elements, unique_elem_range):
+      return dataset_ops.Dataset.range(num_elements).map(
+          lambda x: x % unique_elem_range).unique()
+
+    verify_fn(self, lambda: build_dataset(200, 100), num_outputs=100)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/window_test.py

@@ -0,0 +1,271 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.window()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class WindowTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              count=20,
+              size=[10, 14, 17],
+              shift=[7, 14],
+              stride=[1, 2, 6],
+              drop_remainder=[True, False]) + combinations.combine(
+                  count=[0, 1],
+                  size=10,
+                  shift=4,
+                  stride=1,
+                  drop_remainder=[True, False])))
+  def testWindowDataset(self, count, size, shift, stride, drop_remainder):
+    """Tests a dataset that slides a window its input elements."""
+    components = (np.arange(7),
+                  np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis],
+                  np.array(37.0) * np.arange(7))
+
+    def _map_fn(x, y, z):
+      return math_ops.square(x), math_ops.square(y), math_ops.square(z)
+
+    def _flat_map_fn(x, y, z):
+      return dataset_ops.Dataset.zip((x.batch(batch_size=size),
+                                      y.batch(batch_size=size),
+                                      z.batch(batch_size=size)))
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(components).map(
+        _map_fn).repeat(count).window(
+            size=size,
+            shift=shift,
+            stride=stride,
+            drop_remainder=drop_remainder).flat_map(_flat_map_fn)
+    get_next = self.getNext(dataset)
+
+    self.assertEqual([[None] + list(c.shape[1:]) for c in components],
+                     [ts.as_list() for ts in nest.flatten(
+                         dataset_ops.get_legacy_output_shapes(dataset))])
+
+    num_full_batches = max(0,
+                           (count * 7 - ((size - 1) * stride + 1)) // shift + 1)
+    for i in range(num_full_batches):
+      result = self.evaluate(get_next())
+      for component, result_component in zip(components, result):
+        for j in range(size):
+          self.assertAllEqual(component[(i * shift + j * stride) % 7]**2,
+                              result_component[j])
+    if not drop_remainder:
+      num_partial_batches = (count * 7) // shift + (
+          (count * 7) % shift > 0) - num_full_batches
+      for i in range(num_partial_batches):
+        result = self.evaluate(get_next())
+        for component, result_component in zip(components, result):
+          remaining = (count * 7) - ((num_full_batches + i) * shift)
+          num_elements = remaining // stride + ((remaining % stride) > 0)
+          for j in range(num_elements):
+            self.assertAllEqual(
+                component[((num_full_batches + i) * shift + j * stride) % 7]**2,
+                result_component[j])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(count=20, size=0, shift=3, stride=1) +
+          combinations.combine(count=20, size=3, shift=0, stride=1) +
+          combinations.combine(count=20, size=3, shift=3, stride=0)))
+  def testWindowDatasetInvalid(self, count, size, shift, stride):
+    with self.assertRaises(errors.InvalidArgumentError):
+      ds = dataset_ops.Dataset.range(10).map(lambda x: x).repeat(count).window(
+          size=size, shift=shift,
+          stride=stride).flat_map(lambda x: x.batch(batch_size=size))
+      self.evaluate(ds._variant_tensor)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowDifferentNestedStructures(self):
+    ds = dataset_ops.Dataset.from_tensor_slices(([1, 2], [3, 4])).window(2)
+    self.getNext(ds)
+    ds = dataset_ops.Dataset.from_tensor_slices({"a": [1, 2]}).window(2)
+    self.getNext(ds)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowSparse(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).window(
+        size=5, shift=3,
+        drop_remainder=True).flat_map(lambda x: x.batch(batch_size=5))
+
+    num_batches = (10 - 5) // 3 + 1
+    expected_output = [
+        sparse_tensor.SparseTensorValue(
+            indices=[[0, 0], [1, 0], [2, 0], [3, 0], [4, 0]],
+            values=[i * 3, i * 3 + 1, i * 3 + 2, i * 3 + 3, i * 3 + 4],
+            dense_shape=[5, 1]) for i in range(num_batches)
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowSparseWithDifferentDenseShapes(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=array_ops.expand_dims(
+              math_ops.range(i, dtype=dtypes.int64), 1),
+          values=array_ops.fill([math_ops.cast(i, dtypes.int32)], i),
+          dense_shape=[i])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).window(
+        size=5, shift=3,
+        drop_remainder=True).flat_map(lambda x: x.batch(batch_size=5))
+
+    expected_output = []
+    num_batches = (10 - 5) // 3 + 1
+    for i in range(num_batches):
+      expected_indices = []
+      expected_values = []
+      for j in range(5):
+        for k in range(i * 3 + j):
+          expected_indices.append([j, k])
+          expected_values.append(i * 3 + j)
+      expected_output.append(
+          sparse_tensor.SparseTensorValue(
+              indices=expected_indices,
+              values=expected_values,
+              dense_shape=[5, i * 3 + 5 - 1]))
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedWindowSparse(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+    dataset = dataset_ops.Dataset.range(10).map(_sparse).window(
+        size=4, shift=2,
+        drop_remainder=True).flat_map(lambda x: x.batch(batch_size=4)).window(
+            size=3, shift=1,
+            drop_remainder=True).flat_map(lambda x: x.batch(batch_size=3))
+
+    expected_output = [
+        sparse_tensor.SparseTensorValue(
+            indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], [1, 0, 0],
+                     [1, 1, 0], [1, 2, 0], [1, 3, 0], [2, 0, 0], [2, 1, 0],
+                     [2, 2, 0], [2, 3, 0]],
+            values=[0, 1, 2, 3, 2, 3, 4, 5, 4, 5, 6, 7],
+            dense_shape=[3, 4, 1]),
+        sparse_tensor.SparseTensorValue(
+            indices=[[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 3, 0], [1, 0, 0],
+                     [1, 1, 0], [1, 2, 0], [1, 3, 0], [2, 0, 0], [2, 1, 0],
+                     [2, 2, 0], [2, 3, 0]],
+            values=[2, 3, 4, 5, 4, 5, 6, 7, 6, 7, 8, 9],
+            dense_shape=[3, 4, 1])
+    ]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowShapeError(self):
+
+    def generator():
+      yield [1.0, 2.0, 3.0]
+      yield [4.0, 5.0, 6.0]
+      yield [7.0, 8.0, 9.0, 10.0]
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, dtypes.float32, output_shapes=[None]).window(
+            size=3, shift=1).flat_map(lambda x: x.batch(batch_size=3))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(
+            errors.InvalidArgumentError,
+            r"Cannot batch tensors with different shapes in component 0. "
+            r"First element had shape \[3\] and element 2 had shape \[4\]."))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWindowIgnoreErrors(self):
+    input_values = np.float32([1., np.nan, 2., np.nan, 3.])
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_values).map(
+        lambda x: array_ops.check_numerics(x, "message")).window(
+            size=2, shift=2, stride=2,
+            drop_remainder=True).flat_map(lambda x: x.batch(batch_size=2))
+    self.assertDatasetProduces(
+        dataset, expected_output=[np.float32([1., 2.]),
+                                  np.float32([2., 3.])])
+
+  # Eager-only because the test enumerates the dataset.
+  @combinations.generate(test_base.eager_only_combinations())
+  def testNestedOutput(self):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset_ops.Dataset.zip((dataset, dataset)).window(10)
+    for i, nested_dataset in enumerate(dataset):
+      x, y = nested_dataset
+      self.assertDatasetProduces(x, range(i*10, (i+1)*10))
+      self.assertDatasetProduces(y, range(i*10, (i+1)*10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDropRemainderOutput(self):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset.window(30, drop_remainder=True)
+    dataset = dataset.flat_map(lambda x: x.batch(30))
+    dataset = dataset.batch(4)
+
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[[[y + 30 * x for y in range(30)] for x in range(3)]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = dataset_ops.Dataset.from_tensors(42).window(
+        1, name="window").flat_map(lambda x: x)
+    self.assertDatasetProduces(dataset, [42])
+
+
+class WindowCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                           parameterized.TestCase):
+
+  def _build_dataset(self):
+    dataset = dataset_ops.Dataset.range(42).window(6).interleave(
+        lambda x: x, cycle_length=2, num_parallel_calls=2)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    verify_fn(self, self._build_dataset, num_outputs=42)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/kernel_tests/zip_test.py

@@ -0,0 +1,524 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.data.Dataset.zip()`."""
+import collections
+import dataclasses
+from typing import Callable, Tuple
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import random_access
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+
+try:
+  import attr  # pylint:disable=g-import-not-at-top
+except ImportError:
+  attr = None
+
+
+def _dataset_factory(components):
+  datasets = tuple([
+      dataset_ops.Dataset.from_tensor_slices(component)
+      for component in components
+  ])
+  return dataset_ops.Dataset.zip(datasets)
+
+
+@dataclasses.dataclass
+class MaskedNdarrayPair:
+  mask: bool
+  value1: np.ndarray
+  value2: np.ndarray
+
+  def __tf_flatten__(self):
+    metadata = (self.mask,)
+    components = (self.value1, self.value2)
+    return metadata, components
+
+  def __tf_unflatten__(self, metadata, components):
+    mask = metadata[0]
+    value1, value2 = components
+    return MaskedNdarrayPair(mask=mask, value1=value1, value2=value2)
+
+
+class ZipTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipEqual(self):
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0])
+    ]
+    get_next = self.getNext(_dataset_factory(components))
+    for i in range(4):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component[i], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipUnequal(self):
+    components = [[1, 2, 3, 4], [1, 2, 3, 4, 5], [1.0, 2.0]]
+    get_next = self.getNext(_dataset_factory(components))
+    for i in range(2):
+      results = self.evaluate(get_next())
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component[i], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNested(self):
+
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0])
+    ]
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(component)
+        for component in components
+    ]
+    dataset = dataset_ops.Dataset.zip((datasets[0], (datasets[1], datasets[2])))
+
+    self.assertEqual(
+        dataset_ops.get_legacy_output_shapes(dataset),
+        (tensor_shape.TensorShape([20]),
+         (tensor_shape.TensorShape([22]), tensor_shape.TensorShape([]))))
+
+    get_next = self.getNext(dataset)
+    for i in range(4):
+      result1, (result2, result3) = self.evaluate(get_next())
+      self.assertAllEqual(components[0][i], result1)
+      self.assertAllEqual(components[1][i], result2)
+      self.assertAllEqual(components[2][i], result3)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNamedTuple(self):
+    Foo = collections.namedtuple("Foo", ["x", "y"])
+    x = Foo(x=dataset_ops.Dataset.range(3), y=dataset_ops.Dataset.range(3, 6))
+    dataset = dataset_ops.Dataset.zip(x)
+    expected = [Foo(x=0, y=3), Foo(x=1, y=4), Foo(x=2, y=5)]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataclass(self):
+    mtp = MaskedNdarrayPair(
+        mask=True,
+        value1=dataset_ops.Dataset.range(3),
+        value2=dataset_ops.Dataset.range(3, 6),
+    )
+    dataset = dataset_ops.Dataset.zip(mtp)
+    expected = [
+        MaskedNdarrayPair(mask=True, value1=0, value2=3),
+        MaskedNdarrayPair(mask=True, value1=1, value2=4),
+        MaskedNdarrayPair(mask=True, value1=2, value2=5),
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAttrs(self):
+    if attr is None:
+      self.skipTest("attr module is not available.")
+
+    @attr.s
+    class Foo:
+      x = attr.ib()
+      y = attr.ib()
+
+    x = Foo(x=dataset_ops.Dataset.range(3), y=dataset_ops.Dataset.range(3, 6))
+    dataset = dataset_ops.Dataset.zip(x)
+    expected = [Foo(x=0, y=3), Foo(x=1, y=4), Foo(x=2, y=5)]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    x = dataset_ops.Dataset.from_tensors(4)
+    y = dataset_ops.Dataset.from_tensors(2)
+    dataset = dataset_ops.Dataset.zip((x, y), name="zip")
+    self.assertDatasetProduces(dataset, [(4, 2)])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations())
+  )
+  def testZipWithArgsAndDataset(self):
+    with self.assertRaisesRegex(
+        TypeError, r"Both `\*args` and `datasets` cannot be set."
+    ):
+      dataset_ops.Dataset.zip(
+          dataset_ops.Dataset.range(1, 4),
+          dataset_ops.Dataset.range(4, 7),
+          datasets=(
+              dataset_ops.Dataset.range(1, 4),
+              dataset_ops.Dataset.range(4, 7),
+          ),
+      )
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations())
+  )
+  def testZipBasicWithNoInput(self):
+    with self.assertRaisesRegex(
+        TypeError, r"Must pass at least one dataset to `zip`."
+    ):
+      dataset_ops.Dataset.zip()
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations())
+  )
+  def InvalidZipInputList(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        r"Invalid input to `zip`. Inputs are expected to be (nested)"
+        r" structures of `tf.data.Dataset` objects. Python `list` is"
+        r" not supported and you should use `tuple` instead.",
+    ):
+      dataset_ops.Dataset.zip([1, 2, 3], [4, 5, 6])
+
+
+class ZipCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  def _build_dataset(self, arr, options=None):
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array(arr)
+    ]
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(component)
+        for component in components
+    ]
+    dataset = dataset_ops.Dataset.zip((datasets[0], (datasets[1], datasets[2])))
+    if options:
+      dataset = dataset.with_options(options)
+    return dataset
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(elements=[[37.0, 38.0, 39.0, 40.0], [1.0, 2.0]]),
+          combinations.combine(symbolic_checkpoint=[False, True])))
+  def test(self, verify_fn, elements, symbolic_checkpoint):
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(self, lambda: self._build_dataset(elements, options),
+              len(elements))
+
+
+class ZipRandomAccessTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 3, 4])))
+  def testInvalidIndex(self, index):
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.range(1, 4), dataset_ops.Dataset.range(4, 7)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(index=[-1, 0])))
+  def testEmptyDataset(self, index):
+    dataset = dataset_ops.Dataset.zip(
+        datasets=(dataset_ops.Dataset.from_tensor_slices([]),
+                  dataset_ops.Dataset.from_tensor_slices([])))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=index))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testZipBasic(self):
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.range(1, 4), dataset_ops.Dataset.range(4, 7)))
+    expected_dataset = [(1, 4), (2, 5), (3, 6)]
+    for i in range(3):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)),
+          expected_dataset[i])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testZipBasicWithoutTuple(self):
+    dataset = dataset_ops.Dataset.zip(
+        dataset_ops.Dataset.range(1, 4), dataset_ops.Dataset.range(4, 7)
+    )
+    expected_dataset = [(1, 4), (2, 5), (3, 6)]
+    for i in range(3):
+      self.assertEqual(
+          self.evaluate(random_access.at(dataset, index=i)), expected_dataset[i]
+      )
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations())
+  )
+  def testZipEqual(self):
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0])
+    ]
+    dataset = _dataset_factory(components)
+    for i in range(4):
+      results = self.evaluate(random_access.at(dataset, index=i))
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component[i], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipUnequal(self):
+    components = [[1, 2, 3, 4], [1, 2, 3, 4, 5], [1.0, 2.0]]
+    dataset = _dataset_factory(components)
+    for i in range(2):
+      results = self.evaluate(random_access.at(dataset, index=i))
+      for component, result_component in zip(components, results):
+        self.assertAllEqual(component[i], result_component)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNested(self):
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0])
+    ]
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(component)
+        for component in components
+    ]
+    dataset = dataset_ops.Dataset.zip((datasets[0], (datasets[1], datasets[2])))
+    for i in range(4):
+      result1, (result2,
+                result3) = self.evaluate(random_access.at(dataset, index=i))
+      self.assertAllEqual(components[0][i], result1)
+      self.assertAllEqual(components[1][i], result2)
+      self.assertAllEqual(components[2][i], result3)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedWithoutTuple(self):
+    components = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37.0, 38.0, 39.0, 40.0]),
+    ]
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(component)
+        for component in components
+    ]
+    dataset = dataset_ops.Dataset.zip(datasets[0], (datasets[1], datasets[2]))
+    for i in range(4):
+      result1, (result2, result3) = self.evaluate(
+          random_access.at(dataset, index=i)
+      )
+      self.assertAllEqual(components[0][i], result1)
+      self.assertAllEqual(components[1][i], result2)
+      self.assertAllEqual(components[2][i], result3)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNamedTuple(self):
+    Foo = collections.namedtuple("Foo", ["x", "y"])
+    x = Foo(x=dataset_ops.Dataset.range(3), y=dataset_ops.Dataset.range(3, 6))
+    dataset = dataset_ops.Dataset.zip(x)
+    expected = [Foo(x=0, y=3), Foo(x=1, y=4), Foo(x=2, y=5)]
+    for i in range(3):
+      self.assertAllEqual(
+          self.evaluate(random_access.at(dataset, index=i)), expected[i])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=4))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAttrs(self):
+    if attr is None:
+      self.skipTest("attr module is not available.")
+
+    @attr.s
+    class Foo:
+      x = attr.ib()
+      y = attr.ib()
+
+    x = Foo(x=dataset_ops.Dataset.range(3), y=dataset_ops.Dataset.range(3, 6))
+    dataset = dataset_ops.Dataset.zip(x)
+    expected = [Foo(x=0, y=3), Foo(x=1, y=4), Foo(x=2, y=5)]
+    for i in range(3):
+      self.assertAllEqual(
+          self.evaluate(random_access.at(dataset, index=i)), expected[i])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=4))
+
+
+class ZipGlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          # Tested only on v2 because map v1 does not preserve cardinality.
+          test_base.v2_only_combinations(),
+          combinations.combine(dataset_range=[5, 6, 7]),
+      )
+  )
+  def testZipV2SameLengthInputs(self, dataset_range: int):
+    first_dataset = dataset_ops.Dataset.range(dataset_range)
+    first_dataset = first_dataset.map(lambda x: x * 2)
+    second_dataset = dataset_ops.Dataset.range(dataset_range)
+
+    dataset = dataset_ops.Dataset.zip(first_dataset, second_dataset)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    shuffle_dataset1 = global_shuffle_op._global_shuffle(dataset, seed=10)
+    shuffle_dataset2 = global_shuffle_op._global_shuffle(dataset, seed=11)
+
+    dataset_output1 = self.getDatasetOutput(
+        shuffle_dataset1, requires_initialization=True
+    )
+    dataset_output2 = self.getDatasetOutput(
+        shuffle_dataset2, requires_initialization=True
+    )
+    expected = [(x * 2, x) for x in range(dataset_range)]
+
+    self.assertLen(dataset_output1, dataset_range)
+    self.assertLen(dataset_output2, dataset_range)
+    self.assertCountEqual(dataset_output1, expected)
+    self.assertCountEqual(dataset_output2, expected)
+
+    self.assertNotAllEqual(
+        dataset_output1,
+        dataset_output2,
+        "Different seeds should generate different orders of outputs.",
+    )
+
+    for x_first, x_second in dataset_output1:
+      self.assertEqual(x_first, x_second * 2)
+
+    for x_first, x_second in dataset_output2:
+      self.assertEqual(x_first, x_second * 2)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.v2_only_combinations(),
+          combinations.combine(
+              dataset_ranges=[(10, 8), (9, 5), (4, 7), (5, 8)]
+          ),
+      )
+  )
+  def testZipV2DifferentLengthInputs(self, dataset_ranges: Tuple[int, int]):
+    first_dataset = dataset_ops.Dataset.range(dataset_ranges[0])
+    first_dataset = first_dataset.map(lambda x: x * 2)
+    second_dataset = dataset_ops.Dataset.range(dataset_ranges[1])
+
+    dataset = dataset_ops.Dataset.zip(first_dataset, second_dataset)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    shuffle_dataset1 = global_shuffle_op._global_shuffle(dataset, seed=10)
+    shuffle_dataset2 = global_shuffle_op._global_shuffle(dataset, seed=11)
+
+    dataset_output1 = self.getDatasetOutput(
+        shuffle_dataset1, requires_initialization=True
+    )
+    dataset_output2 = self.getDatasetOutput(
+        shuffle_dataset2, requires_initialization=True
+    )
+    expected = [(x * 2, x) for x in range(min(dataset_ranges))]
+
+    self.assertLen(dataset_output1, min(dataset_ranges))
+    self.assertLen(dataset_output2, min(dataset_ranges))
+    self.assertCountEqual(dataset_output1, expected)
+    self.assertCountEqual(dataset_output2, expected)
+    self.assertNotAllEqual(
+        dataset_output1,
+        dataset_output2,
+        "Different seeds should generate different orders of outputs.",
+    )
+
+    for x_first, x_second in dataset_output1:
+      self.assertEqual(x_first, x_second * 2)
+
+    for x_first, x_second in dataset_output2:
+      self.assertEqual(x_first, x_second * 2)
+
+
+class ZipGlobalShuffleCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.v2_only_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_ranges=[(10, 8), (9, 5), (4, 7), (5, 8)],
+              reshuffle_each_iteration=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testZipV2(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_ranges: Tuple[int, int],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool,
+  ):
+
+    def _build_dataset():
+      first_dataset = dataset_ops.Dataset.range(dataset_ranges[0])
+      first_dataset = first_dataset.map(lambda x: x * 2)
+      second_dataset = dataset_ops.Dataset.range(dataset_ranges[1])
+
+      dataset = dataset_ops.Dataset.zip(first_dataset, second_dataset)
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, reshuffle_each_iteration=reshuffle_each_iteration
+      )
+
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self,
+        _build_dataset,
+        num_outputs=min(dataset_ranges),
+        assert_items_equal=reshuffle_each_iteration,
+    )
+
+
+if __name__ == "__main__":
+  test.main()