chore: import upstream snapshot with attribution

tensorflow/python/data/experimental/BUILD

@@ -0,0 +1,29 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "experimental",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:dataset_ops",
+        "//tensorflow/python/data/experimental/ops:distributed_save_op",
+        "//tensorflow/python/data/experimental/ops:iterator_model_ops",
+        "//tensorflow/python/data/experimental/ops:iterator_ops",
+        "//tensorflow/python/data/experimental/ops:lookup_ops",
+        "//tensorflow/python/data/experimental/ops:parsing_ops",
+        "//tensorflow/python/data/experimental/ops:random_ops",
+        "//tensorflow/python/data/experimental/service",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/data/ops:optional_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/util:all_util",
+    ],
+)

tensorflow/python/data/experimental/__init__.py

@@ -0,0 +1,176 @@
+# 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.
+# ==============================================================================
+"""Experimental API for building input pipelines.
+
+This module contains experimental `Dataset` sources and transformations that can
+be used in conjunction with the `tf.data.Dataset` API. Note that the
+`tf.data.experimental` API is not subject to the same backwards compatibility
+guarantees as `tf.data`, but we will provide deprecation advice in advance of
+removing existing functionality.
+
+See [Importing Data](https://tensorflow.org/guide/datasets) for an overview.
+
+@@AutoShardPolicy
+@@AutotuneAlgorithm
+@@AutotuneOptions
+@@Counter
+@@CsvDataset
+@@DatasetInitializer
+@@DatasetStructure
+@@DistributeOptions
+@@ExternalStatePolicy
+@@OptimizationOptions
+@@Optional
+@@OptionalStructure
+@@RaggedTensorStructure
+@@RandomDataset
+@@Reducer
+@@SparseTensorStructure
+@@SqlDataset
+@@Structure
+@@TFRecordWriter
+@@TensorArrayStructure
+@@TensorStructure
+@@ThreadingOptions
+
+@@assert_cardinality
+@@at
+@@bucket_by_sequence_length
+@@cardinality
+@@choose_from_datasets
+@@copy_to_device
+@@dense_to_ragged_batch
+@@dense_to_sparse_batch
+@@distribute
+@@distributed_save
+@@enable_debug_mode
+@@enumerate_dataset
+@@from_list
+@@from_variant
+@@get_model_proto
+@@get_next_as_optional
+@@get_single_element
+@@get_structure
+@@group_by_reducer
+@@group_by_window
+@@ignore_errors
+@@index_table_from_dataset
+@@load
+@@make_batched_features_dataset
+@@make_csv_dataset
+@@make_saveable_from_iterator
+@@map_and_batch
+@@map_and_batch_with_legacy_function
+@@pad_to_cardinality
+@@parallel_interleave
+@@parse_example_dataset
+@@prefetch_to_device
+@@rejection_resample
+@@sample_from_datasets
+@@save
+@@scan
+@@shuffle_and_repeat
+@@snapshot
+@@table_from_dataset
+@@take_while
+@@to_variant
+@@unbatch
+@@unique
+
+@@AUTOTUNE
+@@INFINITE_CARDINALITY
+@@SHARD_HINT
+@@UNKNOWN_CARDINALITY
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.data.experimental import service
+from tensorflow.python.data.experimental.ops.batching import dense_to_ragged_batch
+from tensorflow.python.data.experimental.ops.batching import dense_to_sparse_batch
+from tensorflow.python.data.experimental.ops.batching import map_and_batch
+from tensorflow.python.data.experimental.ops.batching import map_and_batch_with_legacy_function
+from tensorflow.python.data.experimental.ops.batching import unbatch
+from tensorflow.python.data.experimental.ops.cardinality import assert_cardinality
+from tensorflow.python.data.experimental.ops.cardinality import cardinality
+from tensorflow.python.data.experimental.ops.cardinality import INFINITE as INFINITE_CARDINALITY
+from tensorflow.python.data.experimental.ops.cardinality import UNKNOWN as UNKNOWN_CARDINALITY
+from tensorflow.python.data.experimental.ops.counter import Counter
+from tensorflow.python.data.experimental.ops.distribute import SHARD_HINT
+from tensorflow.python.data.experimental.ops.distributed_save_op import distributed_save
+from tensorflow.python.data.experimental.ops.enumerate_ops import enumerate_dataset
+from tensorflow.python.data.experimental.ops.error_ops import ignore_errors
+from tensorflow.python.data.experimental.ops.from_list import from_list
+from tensorflow.python.data.experimental.ops.get_single_element import get_single_element
+from tensorflow.python.data.experimental.ops.grouping import bucket_by_sequence_length
+from tensorflow.python.data.experimental.ops.grouping import group_by_reducer
+from tensorflow.python.data.experimental.ops.grouping import group_by_window
+from tensorflow.python.data.experimental.ops.grouping import Reducer
+from tensorflow.python.data.experimental.ops.interleave_ops import choose_from_datasets
+from tensorflow.python.data.experimental.ops.interleave_ops import parallel_interleave
+from tensorflow.python.data.experimental.ops.interleave_ops import sample_from_datasets
+from tensorflow.python.data.experimental.ops.io import load
+from tensorflow.python.data.experimental.ops.io import save
+from tensorflow.python.data.experimental.ops.iterator_model_ops import get_model_proto
+from tensorflow.python.data.experimental.ops.iterator_ops import make_saveable_from_iterator
+from tensorflow.python.data.experimental.ops.lookup_ops import DatasetInitializer
+from tensorflow.python.data.experimental.ops.lookup_ops import index_table_from_dataset
+from tensorflow.python.data.experimental.ops.lookup_ops import table_from_dataset
+from tensorflow.python.data.experimental.ops.pad_to_cardinality import pad_to_cardinality
+from tensorflow.python.data.experimental.ops.parsing_ops import parse_example_dataset
+from tensorflow.python.data.experimental.ops.prefetching_ops import copy_to_device
+from tensorflow.python.data.experimental.ops.prefetching_ops import prefetch_to_device
+from tensorflow.python.data.experimental.ops.random_access import at
+from tensorflow.python.data.experimental.ops.random_ops import RandomDataset
+from tensorflow.python.data.experimental.ops.readers import CsvDataset
+from tensorflow.python.data.experimental.ops.readers import make_batched_features_dataset
+from tensorflow.python.data.experimental.ops.readers import make_csv_dataset
+from tensorflow.python.data.experimental.ops.readers import SqlDataset
+from tensorflow.python.data.experimental.ops.resampling import rejection_resample
+from tensorflow.python.data.experimental.ops.scan_ops import scan
+from tensorflow.python.data.experimental.ops.shuffle_ops import shuffle_and_repeat
+from tensorflow.python.data.experimental.ops.snapshot import snapshot
+from tensorflow.python.data.experimental.ops.take_while_ops import take_while
+from tensorflow.python.data.experimental.ops.unique import unique
+from tensorflow.python.data.experimental.ops.writers import TFRecordWriter
+from tensorflow.python.data.ops.dataset_ops import AUTOTUNE
+from tensorflow.python.data.ops.dataset_ops import DatasetSpec as DatasetStructure
+from tensorflow.python.data.ops.dataset_ops import from_variant
+from tensorflow.python.data.ops.dataset_ops import get_structure
+from tensorflow.python.data.ops.dataset_ops import to_variant
+from tensorflow.python.data.ops.debug_mode import enable_debug_mode
+from tensorflow.python.data.ops.iterator_ops import get_next_as_optional
+from tensorflow.python.data.ops.optional_ops import Optional
+from tensorflow.python.data.ops.optional_ops import OptionalSpec as OptionalStructure
+from tensorflow.python.data.ops.options import AutoShardPolicy
+from tensorflow.python.data.ops.options import AutotuneAlgorithm
+from tensorflow.python.data.ops.options import AutotuneOptions
+from tensorflow.python.data.ops.options import DistributeOptions
+from tensorflow.python.data.ops.options import ExternalStatePolicy
+from tensorflow.python.data.ops.options import OptimizationOptions
+from tensorflow.python.data.ops.options import ThreadingOptions
+from tensorflow.python.data.util.structure import _RaggedTensorStructure as RaggedTensorStructure
+from tensorflow.python.data.util.structure import _SparseTensorStructure as SparseTensorStructure
+from tensorflow.python.data.util.structure import _TensorArrayStructure as TensorArrayStructure
+from tensorflow.python.data.util.structure import _TensorStructure as TensorStructure
+from tensorflow.python.framework.type_spec import TypeSpec as Structure
+# pylint: enable=unused-import
+
+from tensorflow.python.util.all_util import remove_undocumented
+
+_allowed_symbols = [
+    "service",
+]
+
+remove_undocumented(__name__, _allowed_symbols)

tensorflow/python/data/experimental/benchmarks/BUILD

@@ -0,0 +1,150 @@
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load(
+    "//tensorflow/tools/test:performance.bzl",
+    "tf_py_benchmark_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_binary(
+    name = "autotune_benchmark_binary",
+    srcs = ["autotune_benchmark.py"],
+    main = "autotune_benchmark.py",
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/ops:math_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "autotune_benchmark",
+    srcs = ["autotune_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/ops:math_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "parameter_value_benchmark",
+    srcs = ["parameter_value_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "csv_dataset_benchmark",
+    srcs = ["csv_dataset_benchmark.py"],
+    tags = ["no_pip"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:test",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "map_and_batch_benchmark",
+    srcs = ["map_and_batch_benchmark.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "map_defun_benchmark",
+    srcs = ["map_defun_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:map_defun",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:map_fn",
+        "//tensorflow/python/ops:math_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "matching_files_benchmark",
+    size = "small",
+    srcs = ["matching_files_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:matching_files",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "optimize_benchmark",
+    srcs = ["optimize_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/ops:math_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "rejection_resample_benchmark",
+    srcs = ["rejection_resample_benchmark.py"],
+    tags = ["no_pip"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:resampling",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "snapshot_dataset_benchmark",
+    srcs = ["snapshot_dataset_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/experimental/ops:snapshot",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "unbatch_benchmark",
+    srcs = ["unbatch_benchmark.py"],
+    deps = [
+        "//tensorflow/python/data/benchmarks:benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)

tensorflow/python/data/experimental/benchmarks/autotune_benchmark.py

@@ -0,0 +1,216 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for autotuning performance knobs."""
+
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.ops import math_ops
+
+
+class AutotuneBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for autotuning performance knobs."""
+
+  def _run_benchmark(self, dataset, autotune, benchmark_iters, benchmark_label,
+                     benchmark_id):
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.autotune.enabled = autotune
+    dataset = dataset.with_options(options)
+
+    autotune_string = "_autotune_parallelism_only"
+    wall_time = self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=benchmark_iters,
+        warmup=True,
+        iters=1,
+        extras={
+            "model_name":
+                "autotune.benchmark.%s.%d" % (benchmark_label, benchmark_id),
+            "parameters":
+                "%s" % autotune,
+        },
+        name=benchmark_label + (autotune_string if autotune else ""))
+    return wall_time
+
+  def benchmark_batch(self):
+    a = self._benchmark_batch(autotune=False, benchmark_id=1)
+    b = self._benchmark_batch(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_batch(self, autotune, benchmark_id):
+    batch_size = 128
+    k = 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(math_ops.matmul)
+    dataset = dataset.batch(
+        batch_size=batch_size, num_parallel_calls=dataset_ops.AUTOTUNE)
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=10000,
+        benchmark_label="batch",
+        benchmark_id=benchmark_id)
+
+  def benchmark_map(self):
+    a = self._benchmark_map(autotune=False, benchmark_id=1)
+    b = self._benchmark_map(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_map(self, autotune, benchmark_id):
+    k = 1024 * 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(
+        math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=10000,
+        benchmark_label="map",
+        benchmark_id=benchmark_id)
+
+  def benchmark_map_and_batch(self):
+    a = self._benchmark_map_and_batch(autotune=False, benchmark_id=1)
+    b = self._benchmark_map_and_batch(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_map_and_batch(self, autotune, benchmark_id):
+    batch_size = 16
+    k = 1024 * 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(
+        math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.batch(batch_size=batch_size)
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=1000,
+        benchmark_label="map_and_batch",
+        benchmark_id=benchmark_id)
+
+  def benchmark_interleave(self):
+    a = self._benchmark_interleave(autotune=False, benchmark_id=1)
+    b = self._benchmark_interleave(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_interleave(self, autotune, benchmark_id):
+    k = 1024 * 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(math_ops.matmul)
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=10000,
+        benchmark_label="interleave",
+        benchmark_id=benchmark_id)
+
+  def benchmark_map_and_interleave(self):
+    a = self._benchmark_map_and_interleave(autotune=False, benchmark_id=1)
+    b = self._benchmark_map_and_interleave(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_map_and_interleave(self, autotune, benchmark_id):
+    k = 1024 * 1024
+    a = (np.random.rand(1, 8 * k), np.random.rand(8 * k, 1))
+    b = (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))
+    c = (np.random.rand(1, 2 * k), np.random.rand(2 * k, 1))
+    dataset_a = dataset_ops.Dataset.from_tensors(a).repeat()
+    dataset_b = dataset_ops.Dataset.from_tensors(b).repeat()
+    dataset_c = dataset_ops.Dataset.from_tensors(c).repeat()
+
+    def f1(x, y):
+      return math_ops.matmul(x, y)
+
+    def f2(a, b):
+      x, y = b
+      return a, math_ops.matmul(x, y)
+
+    dataset = dataset_a
+    dataset = dataset.map(f1, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,
+        num_parallel_calls=dataset_ops.AUTOTUNE,
+        cycle_length=2)
+
+    dataset = dataset_ops.Dataset.zip((dataset, dataset_b))
+    dataset = dataset.map(f2, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,
+        num_parallel_calls=dataset_ops.AUTOTUNE,
+        cycle_length=2)
+
+    dataset = dataset_ops.Dataset.zip((dataset, dataset_c))
+    dataset = dataset.map(f2, num_parallel_calls=dataset_ops.AUTOTUNE)
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=10000,
+        benchmark_label="map_and_interleave",
+        benchmark_id=benchmark_id)
+
+  def benchmark_map_batch_and_interleave(self):
+    a = self._benchmark_map_batch_and_interleave(autotune=False, benchmark_id=1)
+    b = self._benchmark_map_batch_and_interleave(autotune=True, benchmark_id=2)
+    print("autotune parallelism vs no autotuning speedup: {}".format(a / b))
+
+  def _benchmark_map_batch_and_interleave(self, autotune, benchmark_id):
+    batch_size = 16
+    k = 1024 * 1024
+    a = (np.random.rand(1, 8 * k), np.random.rand(8 * k, 1))
+    b = (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))
+    c = (np.random.rand(1, 2 * k), np.random.rand(2 * k, 1))
+    dataset_a = dataset_ops.Dataset.from_tensors(a).repeat()
+    dataset_b = dataset_ops.Dataset.from_tensors(b).repeat()
+    dataset_c = dataset_ops.Dataset.from_tensors(c).repeat()
+
+    dataset = dataset_a
+    dataset = dataset.map(
+        math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.batch(batch_size=batch_size)
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,
+        num_parallel_calls=dataset_ops.AUTOTUNE,
+        cycle_length=2)
+
+    dataset = dataset_ops.Dataset.zip((dataset, dataset_b))
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,
+        num_parallel_calls=dataset_ops.AUTOTUNE,
+        cycle_length=2)
+
+    dataset_c = dataset_c.map(
+        math_ops.matmul, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset_c = dataset_c.batch(batch_size=batch_size)
+    dataset = dataset_ops.Dataset.zip((dataset, dataset_c))
+    return self._run_benchmark(
+        dataset=dataset,
+        autotune=autotune,
+        benchmark_iters=1000,
+        benchmark_label="map_and_batch_and_interleave",
+        benchmark_id=benchmark_id)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/csv_dataset_benchmark.py

@@ -0,0 +1,126 @@
+#  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.
+# ==============================================================================
+"""Benchmarks for `tf.data.experimental.CsvDataset`."""
+
+import os
+import string
+import tempfile
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.ops import readers as core_readers
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import googletest
+
+
+class CsvDatasetBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.experimental.CsvDataset`."""
+
+  FLOAT_VAL = '1.23456E12'
+  STR_VAL = string.ascii_letters * 10
+
+  def _set_up(self, str_val):
+    # Since this isn't test.TestCase, have to manually create a test dir
+    gfile.MakeDirs(googletest.GetTempDir())
+    self._temp_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
+
+    self._num_cols = [4, 64, 256]
+    self._num_per_iter = 5000
+    self._filenames = []
+    for n in self._num_cols:
+      fn = os.path.join(self._temp_dir, 'file%d.csv' % n)
+      with open(fn, 'w') as f:
+        # Just write 100 rows and use `repeat`... Assumes the cost
+        # of creating an iterator is not significant
+        row = ','.join(str_val for _ in range(n))
+        f.write('\n'.join(row for _ in range(100)))
+      self._filenames.append(fn)
+
+  def _tear_down(self):
+    gfile.DeleteRecursively(self._temp_dir)
+
+  def _run_benchmark(self, dataset, num_cols, prefix, benchmark_id):
+
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=self._num_per_iter,
+        name='%s_with_cols_%d' % (prefix, num_cols),
+        iters=10,
+        extras={
+            'model_name': 'csv.benchmark.%d' % benchmark_id,
+            'parameters': '%d' % num_cols,
+        },
+        warmup=True)
+
+  def benchmark_map_with_floats(self):
+    self._set_up(self.FLOAT_VAL)
+    for i in range(len(self._filenames)):
+      num_cols = self._num_cols[i]
+      kwargs = {'record_defaults': [[0.0]] * num_cols}
+      dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
+      dataset = dataset.map(lambda l: parsing_ops.decode_csv(l, **kwargs))  # pylint: disable=cell-var-from-loop
+      self._run_benchmark(
+          dataset=dataset,
+          num_cols=num_cols,
+          prefix='csv_float_map_decode_csv',
+          benchmark_id=1)
+    self._tear_down()
+
+  def benchmark_map_with_strings(self):
+    self._set_up(self.STR_VAL)
+    for i in range(len(self._filenames)):
+      num_cols = self._num_cols[i]
+      kwargs = {'record_defaults': [['']] * num_cols}
+      dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
+      dataset = dataset.map(lambda l: parsing_ops.decode_csv(l, **kwargs))  # pylint: disable=cell-var-from-loop
+      self._run_benchmark(
+          dataset=dataset,
+          num_cols=num_cols,
+          prefix='csv_strings_map_decode_csv',
+          benchmark_id=2)
+    self._tear_down()
+
+  def benchmark_csv_dataset_with_floats(self):
+    self._set_up(self.FLOAT_VAL)
+    for i in range(len(self._filenames)):
+      num_cols = self._num_cols[i]
+      kwargs = {'record_defaults': [[0.0]] * num_cols}
+      dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
+      dataset = readers.CsvDataset(self._filenames[i], **kwargs).repeat()  # pylint: disable=cell-var-from-loop
+      self._run_benchmark(
+          dataset=dataset,
+          num_cols=num_cols,
+          prefix='csv_float_fused_dataset',
+          benchmark_id=3)
+    self._tear_down()
+
+  def benchmark_csv_dataset_with_strings(self):
+    self._set_up(self.STR_VAL)
+    for i in range(len(self._filenames)):
+      num_cols = self._num_cols[i]
+      kwargs = {'record_defaults': [['']] * num_cols}
+      dataset = core_readers.TextLineDataset(self._filenames[i]).repeat()
+      dataset = readers.CsvDataset(self._filenames[i], **kwargs).repeat()  # pylint: disable=cell-var-from-loop
+      self._run_benchmark(
+          dataset=dataset,
+          num_cols=num_cols,
+          prefix='csv_strings_fused_dataset',
+          benchmark_id=4)
+    self._tear_down()
+
+
+if __name__ == '__main__':
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/map_and_batch_benchmark.py

@@ -0,0 +1,199 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.experimental.map_and_batch()`."""
+import hashlib
+import itertools
+
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+
+_NUMPY_RANDOM_SEED = 42
+
+
+class MapAndBatchBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.experimental.map_and_batch()`."""
+
+  def benchmark_map_and_batch(self):
+    """Measures the performance of parallelized batching."""
+    shapes = [(), (10,), (10, 10), (10, 10, 10), (224, 224, 3)]
+    batch_size_values = [1, 32, 64, 128, 1024]
+
+    for shape in shapes:
+      for batch_size in batch_size_values:
+
+        dataset = dataset_ops.Dataset.range(1000000000)
+        dense_value = random_ops.random_normal(shape=shape)
+
+        dataset = dataset.apply(
+            batching.map_and_batch(lambda _: dense_value, batch_size))  # pylint: disable=cell-var-from-loop
+        options = options_lib.Options()
+        options.experimental_optimization.apply_default_optimizations = False
+        dataset = dataset.with_options(options)
+
+        self.run_and_report_benchmark(
+            dataset=dataset,
+            num_elements=batch_size,
+            iters=100,
+            warmup=True,
+            extras={
+                "model_name": "map_and_batch.benchmark.1",
+                "parameters": "%d.%s" % (batch_size, str(shape))
+            },
+            name="num_elements_%d_batch_size_%d" % (np.prod(shape), batch_size))
+
+  def _benchmark_series(self, label, series, benchmark_id):
+    """Runs benchmark the given series."""
+
+    # Decides a proper number of iterations according to the inputs.
+    def compute_num_iters(map_num_calls, inter_op, element_size, batch_size):
+      return 1024 // (
+          (element_size * batch_size) //
+          min(12 if map_num_calls == dataset_ops.AUTOTUNE else map_num_calls,
+              inter_op))
+
+    # Makes the dataset based on the inputs.
+    def make_dataset(map_num_calls, element_size, batch_size, batch_num_calls,
+                     apply_fusion):
+      k = 1024 * 1024
+      x = constant_op.constant(np.random.rand(element_size, 4 * k))
+      y = constant_op.constant(np.random.rand(4 * k, 1))
+      dataset = dataset_ops.Dataset.range(1000000000000).map(lambda _: (x, y))
+      dataset = dataset.map(math_ops.matmul, num_parallel_calls=map_num_calls)
+      dataset = dataset.batch(
+          batch_size=batch_size, num_parallel_calls=batch_num_calls)
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_optimization.map_and_batch_fusion = apply_fusion
+      dataset = dataset.with_options(options)
+      return dataset
+
+    # Makes the name of the dataset based on the inputs.
+    def make_name(label, map_num_calls, inter_op, element_size, batch_size,
+                  batch_num_calls, apply_fusion):
+      map_num_calls_str = ("autotuned" if map_num_calls == dataset_ops.AUTOTUNE
+                           else str(map_num_calls))
+      batch_num_calls_str = (
+          "autotuned" if batch_num_calls == dataset_ops.AUTOTUNE else
+          str(1 if batch_num_calls is None else batch_num_calls))
+      name_str = ("%s_id_%s_map_num_calls_%s_batch_num_calls_%s_inter_op_%d"
+                  "_elem_size_%d_batch_size_%d")
+      name = (
+          name_str % (
+              "fused" if apply_fusion else "chained",
+              hashlib.sha1((label).encode("utf-8")).hexdigest()[:8],
+              map_num_calls_str,
+              batch_num_calls_str,
+              inter_op,
+              element_size,
+              batch_size,
+          ))
+      return name
+
+    for (map_num_calls, inter_op, element_size, batch_size, batch_num_calls,
+         apply_fusion) in series:
+      num_iters = compute_num_iters(map_num_calls, inter_op, element_size,
+                                    batch_size)
+      dataset = make_dataset(map_num_calls, element_size, batch_size,
+                             batch_num_calls, apply_fusion)
+      name = make_name(label, map_num_calls, inter_op, element_size, batch_size,
+                       batch_num_calls, apply_fusion)
+
+      session_config = config_pb2.ConfigProto(
+          inter_op_parallelism_threads=inter_op, use_per_session_threads=True)
+
+      self.run_and_report_benchmark(
+          dataset=dataset,
+          iters=num_iters,
+          num_elements=batch_size,
+          warmup=True,
+          extras={
+              "model_name":
+                  "map_and_batch.benchmark.%d" % benchmark_id,
+              "parameters":
+                  "%d.%d.%d.%d.%d.%s" %
+                  (map_num_calls, inter_op, element_size, batch_size,
+                   batch_num_calls, apply_fusion),
+          },
+          session_config=session_config,
+          name=name)
+
+  def benchmark_map_and_batch_chaining_versus_fusing(self):
+    """Compares the performance of chaining and fusing map and batch.
+
+    NOTE: It is recommended to build the benchmark with
+    `-c opt --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-gmlt`
+    and execute it on a machine with at least 32 CPU cores.
+    """
+
+    # Sequential pipeline configurations.
+    seq_elem_size_series = itertools.product([1], [1], [1, 2, 4, 8], [16],
+                                             [None], [False, True])
+    seq_batch_size_series = itertools.product([1], [1], [1], [8, 16, 32, 64],
+                                              [None], [False, True])
+
+    # Parallel pipeline configuration.
+    par_elem_size_series = itertools.product([32], [32], [1, 2, 4, 8], [256],
+                                             [None], [False, True])
+    par_batch_size_series = itertools.product([32], [32], [1],
+                                              [128, 256, 512, 1024], [None],
+                                              [False, True])
+    par_map_num_calls_series = itertools.product([8, 16, 32, 64], [32], [1],
+                                                 [512], [None], [False, True])
+    par_inter_op_series = itertools.product([32], [8, 16, 32, 64], [1], [512],
+                                            [None], [False, True])
+
+    # Autotuned pipeline configuration.
+    fused_versus_chained_series = [
+        (dataset_ops.AUTOTUNE, 32, 1, 16, dataset_ops.AUTOTUNE, False),
+        (dataset_ops.AUTOTUNE, 32, 1, 16, None, True)
+    ]
+
+    np.random.seed(_NUMPY_RANDOM_SEED)
+    self._benchmark_series(
+        "Sequential element size evaluation",
+        seq_elem_size_series,
+        benchmark_id=2)
+    self._benchmark_series(
+        "Sequential batch size evaluation",
+        seq_batch_size_series,
+        benchmark_id=3)
+    self._benchmark_series(
+        "Parallel element size evaluation",
+        par_elem_size_series,
+        benchmark_id=4)
+    self._benchmark_series(
+        "Parallel batch size evaluation", par_batch_size_series, benchmark_id=5)
+    self._benchmark_series(
+        "Transformation parallelism evaluation",
+        par_map_num_calls_series,
+        benchmark_id=6)
+    self._benchmark_series(
+        "Threadpool size evaluation", par_inter_op_series, benchmark_id=7)
+    self._benchmark_series(
+        "Autotune chained versus fused evaluation",
+        fused_versus_chained_series,
+        benchmark_id=8)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/map_defun_benchmark.py

@@ -0,0 +1,75 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for MapDefunOp."""
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import map_defun
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+
+
+class MapDefunBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for MapDefunOp."""
+
+  def _run(self, op, name, num_iters, benchmark_id):
+
+    wall_time = self.run_op_benchmark(op=op, iters=num_iters, warmup=True)
+    zero_division_delta = 1e-100
+    wall_time = wall_time + zero_division_delta
+    self.report_benchmark(
+        name=name,
+        iters=num_iters,
+        wall_time=wall_time,
+        extras={
+            "examples_per_sec": 1 / float(wall_time),
+            "model_name": "map_defun.benchmark.%d" % benchmark_id,
+            "parameters": "%d" % num_iters,
+        })
+
+  def benchmark_defun_vs_map_fn(self):
+    """Benchmarks to compare the performance of MapDefun vs tf.map_fn."""
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def defun(x):
+      return array_ops.identity(x)
+
+    def fn(x):
+      return array_ops.identity(x)
+
+    base = math_ops.range(10000)
+    for input_size in [10, 100, 1000, 10000]:
+      num_iters = 10000 // input_size
+      map_defun_op = map_defun.map_defun(defun, [base], [dtypes.int32], [()])
+      map_fn_op = map_fn.map_fn(fn, base)
+
+      self._run(
+          op=map_defun_op,
+          name="with_defun_size_%d" % input_size,
+          num_iters=num_iters,
+          benchmark_id=1)
+      self._run(
+          op=map_fn_op,
+          name="without_defun_size_%d" % input_size,
+          num_iters=num_iters,
+          benchmark_id=2)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/matching_files_benchmark.py

@@ -0,0 +1,67 @@
+# 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.
+# ==============================================================================
+"""Benchmark for the experimental `MatchingFilesDataset`."""
+import os
+import shutil
+import tempfile
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import matching_files
+
+
+class MatchingFilesBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmark for the experimental `MatchingFilesDataset`."""
+
+  def benchmark_nested_directories(self):
+    tmp_dir = tempfile.mkdtemp()
+    width = 500
+    depth = 10
+    for i in range(width):
+      for j in range(depth):
+        new_base = os.path.join(tmp_dir, str(i),
+                                *[str(dir_name) for dir_name in range(j)])
+        os.makedirs(new_base)
+        child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
+        for f in child_files:
+          filename = os.path.join(new_base, f)
+          open(filename, 'w').close()
+
+    patterns = [
+        os.path.join(tmp_dir, os.path.join(*['**'
+                                             for _ in range(depth)]), suffix)
+        for suffix in ['*.txt', '*.log']
+    ]
+    # the num_elements depends on the pattern that has been defined above.
+    # In the current scenario, the num of files are selected based on the
+    # ['*.txt', '*.log'] patterns. Since the files which match either of these
+    # patterns are created once per `width`. The num_elements would be:
+    num_elements = width * 2
+
+    dataset = matching_files.MatchingFilesDataset(patterns)
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        iters=3,
+        num_elements=num_elements,
+        extras={
+            'model_name': 'matching_files.benchmark.1',
+            'parameters': '%d.%d' % (width, depth),
+        },
+        name='nested_directory(%d*%d)' % (width, depth))
+
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+if __name__ == '__main__':
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/optimize_benchmark.py

@@ -0,0 +1,165 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for static optimizations."""
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.ops import math_ops
+
+
+class OptimizationBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for static optimizations."""
+
+  def benchmark_map_fusion(self):
+    """Evaluates performance map of fusion."""
+
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      self._benchmark_map_fusion(
+          chain_length=chain_length, optimize_dataset=False)
+      self._benchmark_map_fusion(
+          chain_length=chain_length, optimize_dataset=True)
+
+  def _benchmark_map_fusion(self, chain_length, optimize_dataset):
+
+    dataset = dataset_ops.Dataset.from_tensors(0).repeat(None)
+    for _ in range(chain_length):
+      dataset = dataset.map(lambda x: x)
+    if optimize_dataset:
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_optimization.map_fusion = True
+      dataset = dataset.with_options(options)
+
+    opt_mark = "opt" if optimize_dataset else "noopt"
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=100,
+        iters=10,
+        warmup=True,
+        extras={
+            "model_name": "optimize.benchmark.1",
+            "parameters": "%d.%s" % (chain_length, optimize_dataset),
+        },
+        name="map_fusion_{}_chain_length_{}".format(opt_mark, chain_length))
+
+  def benchmark_map_and_filter_fusion(self):
+    """Evaluates performance map of fusion."""
+
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      self._benchmark_map_and_filter_fusion(
+          chain_length=chain_length, optimize_dataset=False)
+      self._benchmark_map_and_filter_fusion(
+          chain_length=chain_length, optimize_dataset=True)
+
+  def _benchmark_map_and_filter_fusion(self, chain_length, optimize_dataset):
+
+    dataset = dataset_ops.Dataset.from_tensors(0).repeat(None)
+    for _ in range(chain_length):
+      dataset = dataset.map(lambda x: x + 5).filter(
+          lambda x: math_ops.greater_equal(x - 5, 0))
+    if optimize_dataset:
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_optimization.map_and_filter_fusion = True
+      dataset = dataset.with_options(options)
+
+    opt_mark = "opt" if optimize_dataset else "noopt"
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=100,
+        iters=10,
+        warmup=True,
+        extras={
+            "model_name": "optimize.benchmark.2",
+            "parameters": "%d.%s" % (chain_length, optimize_dataset),
+        },
+        name="map_and_filter_fusion_{}_chain_length_{}".format(
+            opt_mark, chain_length))
+
+  # This benchmark compares the performance of pipeline with multiple chained
+  # filter with and without filter fusion.
+
+  def benchmark_filter_fusion(self):
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      self._benchmark_filter_fusion(
+          chain_length=chain_length, optimize_dataset=False)
+      self._benchmark_filter_fusion(
+          chain_length=chain_length, optimize_dataset=True)
+
+  def _benchmark_filter_fusion(self, chain_length, optimize_dataset):
+
+    dataset = dataset_ops.Dataset.from_tensors(5).repeat(None)
+    for _ in range(chain_length):
+      dataset = dataset.filter(lambda x: math_ops.greater_equal(x - 5, 0))
+    if optimize_dataset:
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_optimization.filter_fusion = True
+      dataset = dataset.with_options(options)
+
+    opt_mark = "opt" if optimize_dataset else "noopt"
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=100,
+        iters=10,
+        warmup=True,
+        extras={
+            "model_name": "optimize.benchmark.3",
+            "parameters": "%d.%s" % (chain_length, optimize_dataset),
+        },
+        name="filter_fusion_{}_chain_length_{}".format(opt_mark, chain_length))
+
+  # This benchmark compares the performance of pipeline with multiple chained
+  # filter with and without filter parallelization.
+
+  def benchmark_filter_parallelization(self):
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      self._benchmark_filter_parallelization(
+          chain_length=chain_length, optimize_dataset=False)
+      self._benchmark_filter_parallelization(
+          chain_length=chain_length, optimize_dataset=True)
+
+  def _benchmark_filter_parallelization(self, chain_length, optimize_dataset):
+
+    dataset = dataset_ops.Dataset.from_tensors(5).repeat()
+    for _ in range(chain_length):
+      dataset = dataset.filter(lambda x: math_ops.greater_equal(x - 5, 0))
+    if optimize_dataset:
+      options = options_lib.Options()
+      options.experimental_optimization.apply_default_optimizations = False
+      options.experimental_optimization.filter_parallelization = True
+      dataset = dataset.with_options(options)
+
+    opt_mark = "opt" if optimize_dataset else "noopt"
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=100,
+        iters=10,
+        warmup=True,
+        extras={
+            "model_name": "optimize.benchmark.4",
+            "parameters": "%d.%s" % (chain_length, optimize_dataset),
+        },
+        name="filter_parallelization_{}_chain_length_{}".format(opt_mark,
+                                                                chain_length))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/parameter_value_benchmark.py

@@ -0,0 +1,163 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Benchmarks to compare effect of different parameter values on the performance."""
+import time
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+
+# The maximum sleeping time for each output step and the input sleeping time in
+# milliseconds.
+max_output_sleep_ms = 0.5
+input_sleep_ms = 1.5
+
+
+def sleep_function(x):
+  time.sleep(np.random.uniform(max_output_sleep_ms) / 1000)
+  return x
+
+
+def map_function(x):
+  return script_ops.py_func(sleep_function, [x], x.dtype)
+
+
+class ParameterValueBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks to compare effect of different parameter values on the performance."""
+
+  def _benchmark_map(self, num_parallel_calls, buffer_size):
+    k = 1024 * 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(
+        math_ops.matmul, num_parallel_calls=num_parallel_calls)
+    dataset = dataset.map(map_function)
+    dataset = dataset.prefetch(buffer_size=buffer_size)
+    dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
+
+    name_str = ("map_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f_"
+                "num_parallel_calls_%d_buffer_size_%d")
+    return self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=10000,
+        name=name_str %
+        (max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
+
+  def benchmark_map(self):
+    nums_parallel_calls = [4, 8, 12]
+    buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
+
+    parameters_list = []
+    wall_time_map = {}
+
+    for num_parallel_calls in nums_parallel_calls:
+      for buffer_size in buffer_sizes:
+        parameters = (num_parallel_calls, buffer_size)
+        parameters_list.append(parameters)
+        wall_time = self._benchmark_map(num_parallel_calls, buffer_size)
+        wall_time_map[parameters] = wall_time
+
+    parameters_list.sort(key=lambda x: wall_time_map[x])
+    for parameters in parameters_list:
+      print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
+            wall_time_map[parameters])
+
+  def _benchmark_map_and_batch(self, num_parallel_calls, buffer_size):
+    batch_size = 16
+    k = 1024 * 1024
+
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(
+        math_ops.matmul, num_parallel_calls=num_parallel_calls)
+    dataset = dataset.batch(batch_size=batch_size)
+    dataset = dataset.map(map_function)
+    dataset = dataset.prefetch(buffer_size=buffer_size)
+    dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
+
+    name_str = ("map_and_batch_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f"
+                "_num_parallel_calls_%d_buffer_size_%d")
+    return self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=1000,
+        name=name_str %
+        (max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
+
+  def benchmark_map_and_batch(self):
+    nums_parallel_calls = [4, 8, 12]
+    buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
+
+    parameters_list = []
+    wall_time_map = {}
+
+    for num_parallel_calls in nums_parallel_calls:
+      for buffer_size in buffer_sizes:
+        parameters = (num_parallel_calls, buffer_size)
+        parameters_list.append(parameters)
+        wall_time = self._benchmark_map_and_batch(num_parallel_calls,
+                                                  buffer_size)
+        wall_time_map[parameters] = wall_time
+
+    parameters_list.sort(key=lambda x: wall_time_map[x])
+    for parameters in parameters_list:
+      print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
+            wall_time_map[parameters])
+
+  def _benchmark_interleave(self, num_parallel_calls, buffer_size):
+    k = 1024 * 1024
+    dataset = dataset_ops.Dataset.from_tensors(
+        (np.random.rand(1, 4 * k), np.random.rand(4 * k, 1))).repeat()
+    dataset = dataset.map(math_ops.matmul)
+    dataset = dataset.map(map_function)
+    dataset = dataset_ops.Dataset.range(1).repeat().interleave(
+        lambda _: dataset,  # pylint: disable=cell-var-from-loop
+        cycle_length=10,
+        num_parallel_calls=num_parallel_calls)
+    dataset = dataset.prefetch(buffer_size=buffer_size)
+    dataset = dataset.apply(testing.sleep(int(input_sleep_ms * 1000)))
+
+    name_str = ("interleave_max_output_sleep_ms_%.2f_input_sleep_ms_%.2f"
+                "_num_parallel_calls_%d_buffer_size_%d")
+    return self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=10000,
+        name=name_str %
+        (max_output_sleep_ms, input_sleep_ms, num_parallel_calls, buffer_size))
+
+  def benchmark_interleave(self):
+    nums_parallel_calls = [4, 8, 10]
+    buffer_sizes = [10, 50, 100, 150, 200, 250, 300]
+
+    parameters_list = []
+    wall_time_map = {}
+
+    for num_parallel_calls in nums_parallel_calls:
+      for buffer_size in buffer_sizes:
+        parameters = (num_parallel_calls, buffer_size)
+        parameters_list.append(parameters)
+        wall_time = self._benchmark_interleave(num_parallel_calls, buffer_size)
+        wall_time_map[parameters] = wall_time
+
+    parameters_list.sort(key=lambda x: wall_time_map[x])
+    for parameters in parameters_list:
+      print("num_parallel_calls_%d_buffer_size_%d_wall_time:" % parameters,
+            wall_time_map[parameters])
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/rejection_resample_benchmark.py

@@ -0,0 +1,65 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.experimental.rejection_resample()`."""
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import resampling
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+
+
+class RejectionResampleBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.experimental.rejection_resample()`."""
+
+  def benchmark_resample_performance(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 = 1000
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+    # Prepare the dataset
+    dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat()
+    # Reshape distribution via rejection sampling.
+    dataset = dataset.apply(
+        resampling.rejection_resample(
+            class_func=lambda x: x,
+            target_dist=target_dist,
+            initial_dist=init_dist,
+            seed=142))
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+
+    wall_time = self.run_benchmark(
+        dataset=dataset,
+        num_elements=num_samples,
+        iters=10,
+        warmup=True)
+    resample_time = wall_time * num_samples
+
+    self.report_benchmark(
+        iters=10,
+        wall_time=resample_time,
+        extras={
+            "model_name": "rejection_resample.benchmark.1",
+            "parameters": "%d" % num_samples,
+        },
+        name="resample_{}".format(num_samples))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/snapshot_dataset_benchmark.py

@@ -0,0 +1,200 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.experimental.snapshot()`."""
+import os
+import shutil
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import snapshot
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.platform import test
+
+
+class SnapshotDatasetBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.experimental.snapshot()`."""
+
+  def _makeSnapshotDirectory(self):
+    tmp_dir = test.get_temp_dir()
+    tmp_dir = os.path.join(tmp_dir, "snapshot")
+    if os.path.exists(tmp_dir):
+      shutil.rmtree(tmp_dir)
+    os.mkdir(tmp_dir)
+    return tmp_dir
+
+  def _createSimpleDataset(self,
+                           num_elements,
+                           tmp_dir=None,
+                           compression=snapshot.COMPRESSION_NONE):
+    if not tmp_dir:
+      tmp_dir = self._makeSnapshotDirectory()
+
+    dataset = dataset_ops.Dataset.from_tensor_slices([1.0])
+    dataset = dataset.map(
+        lambda x: gen_array_ops.broadcast_to(x, [50, 50, 3]))
+    dataset = dataset.repeat(num_elements)
+    dataset = dataset.apply(
+        snapshot.legacy_snapshot(tmp_dir, compression=compression))
+
+    return dataset
+
+  def benchmarkWriteSnapshotGzipCompression(self):
+    num_elements = 500000
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements, compression=snapshot.COMPRESSION_GZIP)
+
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="write_gzip",
+        warmup=False,
+        extras={
+            "model_name": "snapshot.benchmark.1",
+            "parameters": "%d" % num_elements,
+        },
+        iters=1)
+
+  def benchmarkWriteSnapshotSnappyCompression(self):
+    num_elements = 500000
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements, compression=snapshot.COMPRESSION_SNAPPY)
+
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="write_snappy",
+        warmup=False,
+        extras={
+            "model_name": "snapshot.benchmark.2",
+            "parameters": "%d" % num_elements,
+        },
+        iters=1)
+
+  def benchmarkWriteSnapshotSimple(self):
+    num_elements = 500000
+    dataset = self._createSimpleDataset(num_elements=num_elements)
+
+    # We only run one iteration here because running multiple iterations will
+    # cause the later iterations to simply read from the already written
+    # snapshot rather than write a new one.
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="write_simple",
+        warmup=False,
+        extras={
+            "model_name": "snapshot.benchmark.3",
+            "parameters": "%d" % num_elements,
+        },
+        iters=1)
+
+  def benchmarkPassthroughSnapshotSimple(self):
+    num_elements = 100000
+    tmp_dir = self._makeSnapshotDirectory()
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements, tmp_dir=tmp_dir)
+
+    # Consume only 1 element, thus making sure we don't finalize.
+    self.run_benchmark(
+        dataset=dataset,
+        num_elements=1,
+        iters=1,
+        warmup=False,
+        apply_default_optimizations=True)
+    # Now run the actual benchmarks and report them
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="passthrough_simple",
+        extras={
+            "model_name": "snapshot.benchmark.4",
+            "parameters": "%d" % num_elements,
+        },
+    )
+
+  def benchmarkReadSnapshotSimple(self):
+    num_elements = 100000
+    tmp_dir = self._makeSnapshotDirectory()
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements, tmp_dir=tmp_dir)
+
+    # consume all the elements to let snapshot write things to disk
+    self.run_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        iters=1,
+        warmup=False,
+        apply_default_optimizations=True)
+    # Now run the actual benchmarks and report them
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="read_simple",
+        extras={
+            "model_name": "snapshot.benchmark.5",
+            "parameters": "%d" % num_elements,
+        })
+
+  def benchmarkReadSnapshotGzipCompression(self):
+    num_elements = 100000
+    tmp_dir = self._makeSnapshotDirectory()
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements,
+        tmp_dir=tmp_dir,
+        compression=snapshot.COMPRESSION_GZIP)
+
+    # consume all the elements to let snapshot write things to disk
+    self.run_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        iters=1,
+        warmup=False,
+        apply_default_optimizations=True)
+    # Now run the actual benchmarks and report them
+    self.run_and_report_benchmark(
+        dataset=dataset, num_elements=num_elements, name="read_gzip",
+        extras={
+            "model_name": "snapshot.benchmark.6",
+            "parameters": "%d" % num_elements,
+        })
+
+  def benchmarkReadSnapshotSnappyCompression(self):
+    num_elements = 100000
+    tmp_dir = self._makeSnapshotDirectory()
+    dataset = self._createSimpleDataset(
+        num_elements=num_elements,
+        tmp_dir=tmp_dir,
+        compression=snapshot.COMPRESSION_SNAPPY)
+
+    # consume all the elements to let snapshot write things to disk
+    self.run_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        iters=1,
+        warmup=False,
+        apply_default_optimizations=True)
+    # Now run the actual benchmarks and report them
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        name="read_snappy",
+        extras={
+            "model_name": "snapshot.benchmark.7",
+            "parameters": "%d" % num_elements,
+        })
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/benchmarks/unbatch_benchmark.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.unbatch()`."""
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+
+
+class UnbatchBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.unbatch()`."""
+
+  def benchmark_native_unbatch(self):
+    batch_sizes = [1, 2, 5, 10, 20, 50]
+    num_elements = 10000
+
+    for batch_size in batch_sizes:
+      dataset = dataset_ops.Dataset.from_tensors("element").repeat(None)
+      dataset = dataset.batch(batch_size)
+      dataset = dataset.unbatch()
+
+      self.run_and_report_benchmark(
+          dataset=dataset,
+          num_elements=num_elements,
+          iters=5,
+          extras={
+              "model_name": "unbatch.benchmark.1",
+              "parameters": "%d" % batch_size,
+          },
+          name="native_batch_size_%d" % batch_size)
+
+  # Include a benchmark of the previous `unbatch()` implementation that uses
+  # a composition of more primitive ops. Eventually we'd hope to generate code
+  # that is as good in both cases.
+  def benchmark_old_unbatch_implementation(self):
+    batch_sizes = [1, 2, 5, 10, 20, 50]
+    num_elements = 10000
+
+    for batch_size in batch_sizes:
+      dataset = dataset_ops.Dataset.from_tensors("element").repeat(None)
+      dataset = dataset.batch(batch_size)
+      dataset = dataset.flat_map(dataset_ops.Dataset.from_tensor_slices)
+
+      self.run_and_report_benchmark(
+          dataset=dataset,
+          num_elements=num_elements,
+          iters=5,
+          extras={
+              "model_name": "unbatch.benchmark.2",
+              "parameters": "%d" % batch_size,
+          },
+          name="unfused_batch_size_%d" % batch_size)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/experimental/kernel_tests/BUILD

@@ -0,0 +1,801 @@
+# Definitions are loaded separately so that copybara can pattern match (and modify) each definition.
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_py_strict_test(
+    name = "assert_cardinality_test",
+    size = "small",
+    srcs = ["assert_cardinality_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:cardinality",
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/experimental/ops:random_access",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "assert_next_test",
+    size = "small",
+    srcs = ["assert_next_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "assert_prev_test",
+    size = "small",
+    srcs = ["assert_prev_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "auto_shard_dataset_test",
+    size = "medium",
+    srcs = ["auto_shard_dataset_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:cardinality",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/experimental/ops:interleave_ops",
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/experimental/ops:unique",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/lib/io:python_io",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "compression_ops_test",
+    size = "small",
+    srcs = ["compression_ops_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:compression_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "copy_to_device_test",
+    size = "small",
+    srcs = ["copy_to_device_test.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:prefetching_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "csv_dataset_test",
+    size = "medium",
+    srcs = ["csv_dataset_test.py"],
+    shard_count = 8,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:error_ops",
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "dense_to_sparse_batch_test",
+    size = "medium",
+    srcs = ["dense_to_sparse_batch_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "from_list_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["from_list_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:from_list",
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/experimental/ops:random_access",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "global_shuffle_test",
+    srcs = ["global_shuffle_test.py"],
+    shard_count = 8,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "group_by_reducer_test",
+    size = "small",
+    srcs = ["group_by_reducer_test.py"],
+    shard_count = 12,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "index_flat_map_test",
+    srcs = ["index_flat_map_test.py"],
+    shard_count = 8,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/experimental/ops:index_flat_map_op",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops/ragged:ragged_string_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "index_shuffle_test",
+    size = "large",
+    srcs = ["index_shuffle_test.py"],
+    shard_count = 12,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:shuffle_ops",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "io_test",
+    size = "medium",
+    srcs = ["io_test.py"],
+    tags = [
+        "notap",  # b/272281090
+    ],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:io",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "lookup_ops_test",
+    size = "small",
+    srcs = ["lookup_ops_test.py"],
+    deps = [
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/data/experimental/ops:lookup_ops",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:lookup_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "make_batched_features_dataset_test",
+    size = "small",
+    srcs = ["make_batched_features_dataset_test.py"],
+    shard_count = 6,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "make_csv_dataset_test",
+    size = "small",
+    srcs = ["make_csv_dataset_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "make_saveable_from_iterator_test",
+    size = "small",
+    srcs = ["make_saveable_from_iterator_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:iterator_ops",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:saver",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "make_tf_record_dataset_test",
+    size = "medium",
+    srcs = ["make_tf_record_dataset_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_and_batch_test",
+    size = "medium",
+    srcs = ["map_and_batch_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/python:pywrap_sanitizers",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_defun_op_test",
+    size = "small",
+    srcs = ["map_defun_op_test.py"],
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/data/experimental/ops:map_defun",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:data_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:sparse_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "matching_files_dataset_test",
+    size = "small",
+    srcs = ["matching_files_dataset_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:matching_files",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "model_dataset_test",
+    size = "small",
+    srcs = ["model_dataset_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "non_serializable_test",
+    size = "small",
+    srcs = ["non_serializable_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "pad_to_cardinality_test",
+    srcs = ["pad_to_cardinality_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:pad_to_cardinality",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "parallel_interleave_test",
+    size = "medium",
+    srcs = ["parallel_interleave_test.py"],
+    shard_count = 8,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:interleave_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/ops:sparse_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "parse_example_dataset_test",
+    size = "medium",
+    srcs = ["parse_example_dataset_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:parsing_ops",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "prefetch_to_device_test",
+    size = "small",
+    srcs = ["prefetch_to_device_test.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:prefetching_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "prefetch_with_slack_test",
+    size = "small",
+    srcs = ["prefetch_with_slack_test.py"],
+    deps = [
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:multi_device_iterator_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "rebatch_dataset_test",
+    size = "medium",
+    srcs = ["rebatch_dataset_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:image_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "replicate_test",
+    size = "medium",
+    srcs = ["replicate_test.py"],
+    grpc_enabled = True,
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "shuffle_and_repeat_test",
+    size = "medium",
+    srcs = ["shuffle_and_repeat_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:shuffle_ops",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "sleep_test",
+    size = "small",
+    srcs = ["sleep_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "sql_dataset_test",
+    size = "medium",
+    srcs = ["sql_dataset_test.py"],
+    shard_count = 4,
+    deps = [
+        "//tensorflow/python/data/experimental/ops:readers",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+        "@org_sqlite//:python",
+    ],
+)
+
+tf_py_strict_test(
+    name = "tf_record_writer_test",
+    size = "medium",
+    srcs = ["tf_record_writer_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/experimental/ops:writers",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/lib/io:python_io",
+        "//tensorflow/python/lib/io:tf_record",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "variant_test",
+    size = "small",
+    srcs = ["variant_test.py"],
+    deps = [
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "weighted_flat_map_test",
+    srcs = ["weighted_flat_map_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:global_shuffle_op",
+        "//tensorflow/python/data/experimental/ops:weighted_flat_map_op",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "wrap_unwrap_test",
+    size = "small",
+    srcs = ["wrap_unwrap_test.py"],
+    deps = [
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "iterator_model_ops_test",
+    size = "small",
+    srcs = ["iterator_model_ops_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:cardinality",
+        "//tensorflow/python/data/experimental/ops:iterator_model_ops",
+        "//tensorflow/python/data/experimental/ops:iterator_ops",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

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

@@ -0,0 +1,155 @@
+# 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.experimental.assert_cardinality()`."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import cardinality
+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 AssertCardinalityTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for `tf.data.experimental.assert_cardinality()`."""
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCorrectCardinality(self):
+    dataset = dataset_ops.Dataset.range(10).filter(lambda x: True)
+    self.assertEqual(
+        self.evaluate(cardinality.cardinality(dataset)), cardinality.UNKNOWN)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+    dataset = dataset.apply(cardinality.assert_cardinality(10))
+    self.assertEqual(self.evaluate(cardinality.cardinality(dataset)), 10)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=20,
+              expected_error="Input dataset was expected to contain 20 "
+              "elements but contained only 10 elements.") +
+          combinations.combine(
+              num_elements=1,
+              asserted_cardinality=20,
+              expected_error="Input dataset was expected to contain 20 "
+              "elements but contained only 1 element.") +
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=cardinality.INFINITE,
+              expected_error="Input dataset was expected to contain an "
+              "infinite number of elements but contained only 10 elements.") +
+          combinations.combine(
+              num_elements=1,
+              asserted_cardinality=cardinality.INFINITE,
+              expected_error="Input dataset was expected to contain an "
+              "infinite number of elements but contained only 1 element.") +
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=5,
+              expected_error="Input dataset was expected to contain 5 "
+              "elements but contained at least 6 elements.") +
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=1,
+              expected_error="Input dataset was expected to contain 1 "
+              "element but contained at least 2 elements.")))
+  def testIncorrectCardinality(self, num_elements, asserted_cardinality,
+                               expected_error):
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.apply(
+        cardinality.assert_cardinality(asserted_cardinality))
+    get_next = self.getNext(dataset)
+    with self.assertRaisesRegex(errors.FailedPreconditionError, expected_error):
+      while True:
+        self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=100,
+              expected_error=errors.FailedPreconditionError,
+              expected_error_message=(
+                  "Input dataset was expected to contain 100 elements.")) +
+          combinations.combine(
+              num_elements=10,
+              asserted_cardinality=cardinality.INFINITE,
+              expected_error=errors.InvalidArgumentError,
+              expected_error_message=(
+                  "`global_shuffle` requires the input dataset to have a "
+                  "non-empty finite cardinality."))))
+  def testIncorrectCardinalityForGlobalShuffle(
+      self,
+      num_elements: int,
+      asserted_cardinality: int,
+      expected_error: Exception,
+      expected_error_message: str):
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.apply(
+        cardinality.assert_cardinality(asserted_cardinality))
+    with self.assertRaisesRegex(
+        expected_error, expected_error_message):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRandomAccess(self):
+    num_elements = 10
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.apply(cardinality.assert_cardinality(num_elements))
+    self.verifyRandomAccess(dataset, expected=range(num_elements))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRandomAccessOutOfRange(self):
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.apply(cardinality.assert_cardinality(10))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(random_access.at(dataset, index=10))
+    with self.assertRaises(errors.FailedPreconditionError):
+      self.evaluate(random_access.at(dataset, index=5))
+
+
+class AssertCardinalityCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                      parameterized.TestCase):
+
+  def build_dataset(self, num_elements, options=None):
+    dataset = dataset_ops.Dataset.range(num_elements).apply(
+        cardinality.assert_cardinality(num_elements))
+    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(200, options), num_outputs=200)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,73 @@
+# 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.experimental.assert_next()`."""
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class AssertNextTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertNext(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_next(["Map"])).map(lambda x: x)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIgnoreVersionSuffix(self):
+    # The `batch` transformation creates a "BatchV2" dataset, but we should
+    # still match that with "Batch".
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_next(["Map", "Batch"])).map(lambda x: x).batch(1)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[[0]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertNextInvalid(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_next(["Whoops"]))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(errors.InvalidArgumentError,
+                        "Asserted transformation matching Whoops"))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertNextShort(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_next(["Root", "Whoops"]))
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(
+            errors.InvalidArgumentError,
+            "Asserted next 2 transformations but encountered only 1."))
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,86 @@
+# 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.experimental.assert_prev()`."""
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class AssertPrevTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertPrev(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).map(
+        lambda x: x, deterministic=True, num_parallel_calls=8).apply(
+            testing.assert_prev([("ParallelMapDataset",
+                                  {"deterministic", "true"})]))
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIgnoreVersionSuffix(self):
+    # The `batch` transformation creates a "BatchV2" dataset, but we should
+    # still match that with "Batch".
+    dataset = dataset_ops.Dataset.from_tensors(0).map(
+        lambda x: x, deterministic=True, num_parallel_calls=8).batch(1).apply(
+            testing.assert_prev([("BatchDataset", {}),
+                                 ("ParallelMapDataset", {
+                                     "deterministic": "true"
+                                 })]))
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[[0]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertPrevInvalid(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_prev([("Whoops", {})]))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(errors.InvalidArgumentError,
+                        "Asserted transformation matching 'Whoops'"))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertPrevShort(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_prev([("TensorDataset", {}), ("Whoops", {})]))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(
+            errors.InvalidArgumentError,
+            "Asserted previous 2 transformations but encountered only 1."))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertBadAttributeName(self):
+    dataset = dataset_ops.Dataset.from_tensors(0).apply(
+        testing.assert_prev([("TensorDataset", {
+            "whoops": "true"
+        })]))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(errors.InvalidArgumentError, "found no such attribute"))
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,762 @@
+# 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 private `_AutoShardDataset` transformation."""
+import os
+from typing import Optional
+
+from absl.testing import parameterized
+
+from tensorflow.core.example import example_pb2
+from tensorflow.core.example import feature_pb2
+from tensorflow.python.data.experimental.ops import cardinality
+from tensorflow.python.data.experimental.ops import distribute
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.experimental.ops import unique
+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 as core_readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+
+
+def chunk(l, n):
+  for i in range(0, len(l), n):
+    yield l[i:i + n]
+
+
+class AutoShardDatasetTest(tf_record_test_base.TFRecordTestBase,
+                           parameterized.TestCase):
+
+  def setUp(self):
+    super(AutoShardDatasetTest, self).setUp()
+    self._num_files = 10
+    self._num_records = 10
+    self._filenames = self._createFiles()
+
+  def getAllDatasetElements(self, dataset):
+    actual = []
+    next_fn = self.getNext(dataset)
+    while True:
+      try:
+        actual.append(self.evaluate(next_fn()))
+      except errors.OutOfRangeError:
+        break
+    return actual
+
+  def assertDatasetProducesWithShuffle(self, dataset, expected, batch,
+                                       num_examples, shuffle):
+    if shuffle:
+      actual = []
+      next_fn = self.getNext(dataset)
+      for _ in range(num_examples):
+        elem = self.evaluate(next_fn())
+        if isinstance(elem, tuple):
+          actual.extend(elem)
+        else:
+          actual.extend(elem.tolist())
+
+      self.assertCountEqual(actual, expected)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_fn())
+    else:
+      self.assertDatasetProduces(dataset, list(chunk(expected, batch)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(shuffle=[True, False])))
+  def testFlatMapReaderPipeline(self, shuffle):
+    dataset = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=shuffle)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in (3, 8)
+        for r in range(0, 10)
+    ]
+    self.assertDatasetProducesWithShuffle(dataset, expected, 5, 4, shuffle)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(batch_size=[1, 3, 10])))
+  def testDatasetOfReaderDatasetsPipeline(self, batch_size):
+    # This tests a scenario where a list_files main return multiple files
+    # due to the glob containing wildcards.
+    def batch(iterator, n):
+      l = len(iterator)
+      for i in range(0, l, n):
+        yield iterator[i:min(i + n, l)]
+
+    datasets = []
+    for files in batch(self._filenames, batch_size):
+      datasets.append(
+          dataset_ops.Dataset.list_files(files, shuffle=False).map(
+              core_readers.TFRecordDataset))
+    dataset = dataset_ops.Dataset.from_tensor_slices(datasets)
+    dataset = dataset.flat_map(lambda x: x)
+
+    # Simulate additional ops in between flat_map and interleave. This should be
+    # a no-op since if ShardDataset is placed right after flat_map, we will only
+    # have two datasets left at this point.
+    dataset = dataset.prefetch(1)
+    dataset = dataset.prefetch(1)
+
+    dataset = dataset.interleave(
+        lambda x: x, cycle_length=1, num_parallel_calls=1)
+
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in (0, 5)
+        for r in range(0, 10)
+    ]
+
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipReaderPipeline(self):
+    dataset1 = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=False)
+    dataset1 = dataset1.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset2 = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=False)
+    dataset2 = dataset2.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        (b"Record %d of file %d" % (r, f), b"Record %d of file %d" % (r, f))  # pylint:disable=g-complex-comprehension
+        for r in range(0, 10)
+        for f in (3, 8)
+    ]
+
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(shuffle=[True, False])))
+  def testConcatenateReaderPipeline(self, shuffle):
+    dataset1 = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=shuffle)
+    dataset1 = dataset1.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset1 = dataset1.batch(5)
+    dataset2 = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=shuffle)
+    dataset2 = dataset2.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset2 = dataset2.batch(5)
+
+    dataset = dataset1.concatenate(dataset2)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for r in range(0, 10)
+        for f in (3, 8)
+    ]
+    expected += expected
+    self.assertDatasetProducesWithShuffle(dataset, expected, 5, 8, shuffle)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(shuffle=[True, False])))
+  def testPipelineWithMap(self, shuffle):
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset = dataset.map(lambda x: string_ops.substr_v2(x, 2, 1000))
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        b"cord %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for r in range(0, 10)
+        for f in (3, 8)
+    ]
+    self.assertDatasetProducesWithShuffle(dataset, expected, 5, 4, shuffle)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDirectFilenameTFRecordReaderPipeline(self):
+    dataset = core_readers.TFRecordDataset(self._filenames)
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in (0, 5)
+        for r in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(shuffle=[True, False])))
+  def testValidPipelineWithRangeDataset(self, shuffle):
+    dataset = dataset_ops.Dataset.range(self._num_files)
+    dataset = dataset.map(lambda n: string_ops.string_join(  # pylint:disable=g-long-lambda
+        [self.get_temp_dir(),
+         string_ops.string_format("/tf_record.{}.txt", [n])]))
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset = dataset.map(lambda x: string_ops.substr_v2(x, 2, 1000))
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        b"cord %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for r in range(0, 10)
+        for f in (3, 8)
+    ]
+    self.assertDatasetProducesWithShuffle(dataset, expected, 5, 4, shuffle)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(params=[(1, 0, 10, 10), (2, 1, 20, 5),
+                                       (10, 1, 1, 10)])))
+  def testStandardReaderPipeline(self, params):
+    num_epochs, index, batch_size, parallel_reads = params
+    dataset = readers.make_tf_record_dataset(
+        file_pattern=self._filenames,
+        num_epochs=num_epochs,
+        batch_size=batch_size,
+        parser_fn=None,
+        num_parallel_reads=parallel_reads,
+        drop_final_batch=True,
+        shuffle=False)
+    dataset = distribute._AutoShardDataset(dataset, 2, index)
+    outputs = self.getNext(dataset)
+    self._verify_records(
+        outputs,
+        batch_size=batch_size,
+        file_index=[i for i in range(index, self._num_records, 2)],
+        num_epochs=num_epochs,
+        interleave_cycle_length=parallel_reads,
+        drop_final_batch=True,
+        use_parser_fn=None)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(outputs())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShardInputToInterleave(self):
+    file1 = self._writeFile("f0", [1, 2, 3])
+    file2 = self._writeFile("f1", [4, 5, 6])
+    file3 = self._writeFile("f2", [7, 8, 9])
+    dataset = dataset_ops.Dataset.from_tensor_slices([file1, file2, file3])
+    dataset = dataset.interleave(core_readers.TFRecordDataset, cycle_length=3)
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+
+    # Sharding by file will interleave files 0 and 2
+    expected = [str.encode(str(i)) for i in [1, 7, 2, 8, 3, 9]]
+    actual = self.getDatasetOutput(dataset)
+    self.assertEqual(actual, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShardInputToInterleaveWithIdentityFunction(self):
+    file1 = self._writeFile("f0", [1, 2, 3])
+    file2 = self._writeFile("f1", [4, 5, 6])
+    file3 = self._writeFile("f2", [7, 8, 9])
+    dataset = dataset_ops.Dataset.from_tensor_slices([file1, file2, file3])
+    dataset = dataset.map(core_readers.TFRecordDataset)
+    dataset = dataset.interleave(lambda x: x, cycle_length=3)
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+
+    # Sharding by file will interleave files 0 and 2
+    expected = [str.encode(str(i)) for i in [1, 7, 2, 8, 3, 9]]
+    actual = self.getDatasetOutput(dataset)
+    self.assertEqual(actual, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(shuffle=[True, False])))
+  def testSampleResNetPipeline(self, shuffle):
+    dataset = dataset_ops.Dataset.list_files(
+        self._filenames, shuffle=shuffle)
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for r in range(0, 10)
+        for f in (3, 8)
+    ]
+    self.assertDatasetProducesWithShuffle(dataset, expected, 5, 4, shuffle)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              options_lib.AutoShardPolicy.DATA,
+              options_lib.AutoShardPolicy.AUTO
+          ])))
+  def testShardByDataBeforePrefetch(self, sharding_policy):
+    dataset = dataset_ops.Dataset.range(4)
+    dataset = dataset.apply(testing.assert_next(["Shard", "Prefetch"]))
+    dataset = dataset.prefetch(1)
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = sharding_policy
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+    self.assertDatasetProduces(dataset, [0, 2])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.times(combinations.combine(
+              sharding_policy=[options_lib.AutoShardPolicy.DATA,
+                               options_lib.AutoShardPolicy.FILE]),
+                             combinations.combine(shuffle=[True, False]))))
+  def testReplicateAndShardProduceDisjointData(self, shuffle, sharding_policy):
+    dataset = dataset_ops.Dataset.list_files(self._filenames,
+                                             shuffle=shuffle)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+
+    graph_def = dataset._as_serialized_graph(
+        strip_device_assignment=True,
+        external_state_policy=options_lib.ExternalStatePolicy.WARN)
+
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = sharding_policy
+
+    ds1 = distribute._RemoteDataset(graph_def, "/device:CPU:0",
+                                    dataset.element_spec)
+    ds2 = distribute._RemoteDataset(graph_def, "/device:CPU:0",
+                                    dataset.element_spec)
+
+    ds1 = ds1.with_options(options)
+    ds2 = ds2.with_options(options)
+
+    ds1 = distribute._AutoShardDataset(ds1, 2, 0)
+    ds2 = distribute._AutoShardDataset(ds2, 2, 1)
+
+    elems1 = set(self.getAllDatasetElements(ds1))
+    elems2 = set(self.getAllDatasetElements(ds2))
+
+    self.assertEmpty(elems1.intersection(elems2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWorkersGreaterThanNumFilesWithDataSharding(self):
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.DATA)
+
+    dataset = core_readers._TFRecordDataset(self._filenames)
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    # Should return "Record (0,5) of file (0 --> 9)" since we are sharding by
+    # individual elements, we should be able to get some data from all files.
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in range(0, 10)
+        for r in (0, 5)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAutoshardPolicyOff(self):
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.OFF)
+
+    dataset = core_readers._TFRecordDataset(self._filenames)
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    # Should return every record in every file since autosharding is turned off.
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in range(0, 10)
+        for r in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileShardingWithoutReaderDatasetOp(self):
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.FILE)
+
+    dataset = dataset_ops.Dataset.range(1024)
+    dataset = dataset.with_options(options)
+
+    # We are specifying that we want a file sharding policy, and this pipeline
+    # doesn't start with file reading, so we should error out.
+    with self.assertRaises(errors.NotFoundError):
+      dataset = distribute._AutoShardDataset(dataset, 10, 0)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWorkersGreaterThanNumFiles(self):
+    dataset = dataset_ops.Dataset.list_files(self._filenames)
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 500, 499)
+    self.assertDatasetProduces(dataset, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordReaderWithDirectFileNames(self):
+    # Using `_TFRecordDataset` creates a raw op rather than wrapping it around
+    # a flat_map automatically.
+    dataset = core_readers._TFRecordDataset(self._filenames)
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in range(0, 10)
+        for r in (0, 5)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordReaderWithDirectFileNamesAndShapes(self):
+    # Using `_TFRecordDataset` creates a raw op rather than wrapping it around
+    # a flat_map automatically.
+    dataset = core_readers._TFRecordDataset(self._filenames)
+
+    # BatchDataset contains `output_types` and `output_shapes`
+    dataset = dataset.batch(5)
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in range(0, 10)
+        for r in range(0, 5)
+    ]
+    self.assertDatasetProduces(dataset, list(chunk(expected, 5)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShardOutOfRange(self):
+    dataset = dataset_ops.Dataset.range(5)
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = distribute._AutoShardDataset(dataset, 10, 0)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShardOutOfRangeEmptyDataset(self):
+    dataset = dataset_ops.Dataset.range(0)
+    with self.assertRaises(errors.OutOfRangeError):
+      dataset = distribute._AutoShardDataset(dataset, 10, 0)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoReaderPipelines(self):
+    dataset = dataset_ops.Dataset.range(1024)
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+    self.assertDatasetProduces(dataset, [i for i in range(1024) if i % 2 == 0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnknownOpInPipelineStillShardsAtTheEnd(self):
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.apply(unique.unique())
+
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in range(0, 10)
+        for r in (0, 5)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidWorkerIndex(self):
+    dataset = dataset_ops.Dataset.list_files(self._filenames)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = distribute._AutoShardDataset(dataset, 2, 2)
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertCardinality(self):
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+    dataset = dataset.apply(cardinality.assert_cardinality(42))
+    dataset = distribute._AutoShardDataset(dataset, 5, 0)
+
+    expected = [
+        b"Record %d of file %d" % (r, f)  # pylint:disable=g-complex-comprehension
+        for f in (0, 5)
+        for r in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, list(chunk(expected, 5)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMakeBatchedFeaturesDataset(self):
+    files = 2
+    records_per_file = 5
+
+    def make_record(file_index):
+      example = example_pb2.Example(
+          features=feature_pb2.Features(
+              feature={
+                  "file":
+                      feature_pb2.Feature(
+                          int64_list=feature_pb2.Int64List(value=[file_index])),
+              }))
+      return example.SerializeToString()
+
+    filenames = []
+    for file_index in range(files):
+      filename = os.path.join(self.get_temp_dir(),
+                              "tf_record.%d.txt" % file_index)
+      filenames.append(filename)
+      writer = python_io.TFRecordWriter(filename)
+      for _ in range(records_per_file):
+        writer.write(make_record(file_index))
+      writer.close()
+
+    dataset = readers.make_batched_features_dataset(
+        file_pattern=filenames,
+        batch_size=records_per_file,
+        features={
+            "file": parsing_ops.FixedLenFeature([], dtypes.int64),
+        },
+        reader=core_readers.TFRecordDataset,
+        num_epochs=1)
+    # We should shard at the file level, so that all records come from file 0.
+    dataset = distribute._AutoShardDataset(dataset, 2, 0)
+    dataset = dataset.unbatch()
+    output = self.getDatasetOutput(dataset)
+    files = [elem["file"] for elem in output]
+    self.assertEqual(files, [0] * records_per_file)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testHintShardingValidPattern(self):
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.HINT)
+
+    dataset = dataset_ops.Dataset.range(100).shard(distribute.SHARD_HINT, 0)
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 10, 0)
+
+    self.assertDatasetProduces(dataset, list(range(0, 100, 10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testHintShardingInvalidPattern(self):
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        options_lib.AutoShardPolicy.HINT)
+
+    dataset = dataset_ops.Dataset.range(100).shard(1, 0)
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 10, 0)
+
+    self.assertDatasetProduces(dataset, list(range(100)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              auto_shard_policy=list(
+                  policy.name for policy in options_lib.AutoShardPolicy
+              )
+          ),
+      )
+  )
+  def testEnumerateAutoShardPolicies(self, auto_shard_policy):
+    """Verifies tf.data handles every auto-shard policy with no errors."""
+    policy_enum = options_lib.AutoShardPolicy[auto_shard_policy]
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = policy_enum
+    dataset = dataset.with_options(options)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+    self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class AutoShardWithRebatchDatasetTest(tf_record_test_base.TFRecordTestBase,
+                                      parameterized.TestCase):
+
+  def _setUpFiles(self, num_files, num_records_per_file):
+    self._num_files = num_files
+    self._num_records = num_records_per_file
+    self._filenames = self._createFiles()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileShardingWithLegacyRebatch(self):
+    # Tests that RebatchDatasetV1 is a passthrough op.
+    self._setUpFiles(num_files=5, num_records_per_file=10)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.apply(
+        testing.assert_next(["Shard", "FlatMap", "Batch", "Rebatch"]))
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+    dataset = distribute._LegacyRebatchDataset(dataset, num_replicas=5)
+    dataset = distribute._AutoShardDataset(dataset, 5, 3)
+    expected = [[self._record(3, i)] for i in range(10)]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileShardingWithRebatch(self):
+    # Tests that RebatchDatasetV2 is a passthrough op.
+    self._setUpFiles(num_files=3, num_records_per_file=5)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.apply(
+        testing.assert_next(["Shard", "FlatMap", "Batch", "Rebatch"]))
+    dataset = dataset.flat_map(core_readers.TFRecordDataset)
+    dataset = dataset.batch(5)
+    dataset = dataset.rebatch(batch_size=[2, 1, 2])
+    dataset = distribute._AutoShardDataset(dataset, 3, 1)
+    expected = [[self._record(1, 0), self._record(1, 1)], [self._record(1, 2)],
+                [self._record(1, 3), self._record(1, 4)]]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.times(
+              combinations.combine(sharding_policy=[
+                  options_lib.AutoShardPolicy.DATA,
+                  options_lib.AutoShardPolicy.AUTO
+              ]), combinations.combine(with_prefetch=[True, False]))))
+  def testUseLegacyRebatchWithDataSharding(self, sharding_policy,
+                                           with_prefetch):
+    # This test simulates a distributed environment with 3 workers, each with
+    # 1 replica.
+    dataset = dataset_ops.Dataset.range(8)
+    dataset = dataset.batch(4)
+    options = options_lib.Options()
+    options.experimental_distribute.auto_shard_policy = sharding_policy
+    dataset = dataset.with_options(options)
+    # We expect the auto-shard rewrite to rewrite RebatchDatasetV2 to
+    # RebatchDataset(V1) for correctness reasons. This will modify the output
+    # of the dataset.
+    worker_a_dataset = dataset.rebatch(batch_size=[2, 1, 1])
+    if with_prefetch:
+      worker_a_dataset = worker_a_dataset.prefetch(1)
+    worker_a_dataset = distribute._AutoShardDataset(
+        worker_a_dataset, 3, 0, num_replicas=3)
+    expected = [[0, 1], [4, 5]]
+    self.assertDatasetProduces(worker_a_dataset, expected)
+
+    worker_b_dataset = dataset.rebatch(batch_size=[1, 1, 2])
+    if with_prefetch:
+      worker_b_dataset = worker_b_dataset.prefetch(1)
+    worker_b_dataset = distribute._AutoShardDataset(
+        worker_b_dataset, 3, 1, num_replicas=3)
+    expected = [[2, 3], [6, 7]]
+    self.assertDatasetProduces(worker_b_dataset, expected)
+
+    worker_c_dataset = dataset.rebatch(batch_size=[1, 2, 1])
+    if with_prefetch:
+      worker_c_dataset = worker_c_dataset.prefetch(1)
+    worker_c_dataset = distribute._AutoShardDataset(
+        worker_c_dataset, 3, 2, num_replicas=3)
+    expected = [[], []]
+    self.assertDatasetProduces(worker_c_dataset, expected)
+
+
+class AutoShardDatasetCheckpointTest(tf_record_test_base.TFRecordTestBase,
+                                     checkpoint_test_base.CheckpointTestBase,
+                                     parameterized.TestCase):
+
+  def setUp(self):
+    super(AutoShardDatasetCheckpointTest, self).setUp()
+    self._num_files = 10
+    self._num_records = 10
+    self._filenames = self._createFiles()
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+
+    def build_dataset():
+      dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+      dataset = dataset.apply(
+          interleave_ops.parallel_interleave(core_readers.TFRecordDataset, 10))
+      dataset = distribute._AutoShardDataset(dataset, 5, 3)
+      return dataset
+
+    verify_fn(self, build_dataset, num_outputs=20)
+
+
+class AutoShardGlobalShuffleTest(
+    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 test(
+      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 = distribute._AutoShardDataset(dataset, num_shards, shard_index)
+    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)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNotSufficientInput(self):
+    dataset = dataset_ops.Dataset.range(1)
+    dataset = distribute._AutoShardDataset(dataset, 5, 4)
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,169 @@
+# 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 compression ops."""
+from collections import namedtuple
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import compression_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+
+
+def _test_objects():
+
+  Item = namedtuple("Item", "id name")
+
+  return [
+      combinations.NamedObject("int", 1),
+      combinations.NamedObject("string", "dog"),
+      combinations.NamedObject("tuple", (1, 1)),
+      combinations.NamedObject("nested_tuple", ((1, 1), (2, 2))),
+      combinations.NamedObject("named_tuple", Item(id=1, name="item1")),
+      combinations.NamedObject("unicode", "アヒル"),
+      combinations.NamedObject(
+          "nested_named_tuple",
+          (Item(id=1, name="item1"), Item(id=2, name="item2"))),
+      combinations.NamedObject("int_string_tuple", (1, "dog")),
+      combinations.NamedObject(
+          "sparse",
+          sparse_tensor.SparseTensorValue(
+              indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])),
+      combinations.NamedObject(
+          "sparse_structured", {
+              "a":
+                  sparse_tensor.SparseTensorValue(
+                      indices=[[0, 0], [1, 2]],
+                      values=[1, 2],
+                      dense_shape=[3, 4]),
+              "b": (1, 2, "dog")
+          })
+  ]
+
+
+def _test_v2_eager_only_objects():
+  return [
+      combinations.NamedObject(
+          "ragged",
+          ragged_factory_ops.constant([[0, 1, 2, 3], [4, 5], [6, 7, 8], [9]])),
+      combinations.NamedObject(
+          "sparse_ragged_structured", {
+              "sparse":
+                  sparse_tensor.SparseTensorValue(
+                      indices=[[0, 0], [1, 2]],
+                      values=[1, 2],
+                      dense_shape=[3, 4]),
+              "ragged":
+                  ragged_factory_ops.constant([[0, 1, 2, 3], [9]])
+          })
+  ]
+
+
+class CompressionOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(element=_test_objects())) +
+      combinations.times(
+          test_base.v2_eager_only_combinations(),
+          combinations.combine(element=_test_v2_eager_only_objects())))
+  def testCompression(self, element):
+    element = element._obj
+
+    compressed = compression_ops.compress(element)
+    uncompressed = compression_ops.uncompress(
+        compressed, structure.type_spec_from_value(element))
+    self.assertValuesEqual(element, self.evaluate(uncompressed))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(element=_test_objects())) +
+      combinations.times(
+          test_base.v2_eager_only_combinations(),
+          combinations.combine(element=_test_v2_eager_only_objects())))
+  def testDatasetCompression(self, element):
+    element = element._obj
+
+    dataset = dataset_ops.Dataset.from_tensors(element)
+    element_spec = dataset.element_spec
+
+    dataset = dataset.map(lambda *x: compression_ops.compress(x))
+    dataset = dataset.map(lambda x: compression_ops.uncompress(x, element_spec))
+    self.assertDatasetProduces(dataset, [element])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testCompressionOutputDTypeMismatch(self):
+    element = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    compressed = compression_ops.compress(element)
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                "but got a tensor of type string"):
+      uncompressed = compression_ops.uncompress(
+          compressed, structure.type_spec_from_value(0))
+      self.evaluate(uncompressed)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testCompressionInputShapeMismatch(self):
+    element = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    compressed = compression_ops.compress(element)
+    compressed = [compressed, compressed]
+    error = (
+        errors.InvalidArgumentError
+        if context.executing_eagerly() else ValueError)
+    with self.assertRaises(error):
+      uncompressed = compression_ops.uncompress(
+          compressed, structure.type_spec_from_value(0))
+      self.evaluate(uncompressed)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testCompressionInputDTypeMismatch(self):
+    uncompressed = list(range(10))
+    with self.assertRaises(TypeError):
+      uncompressed = compression_ops.uncompress(
+          uncompressed, structure.type_spec_from_value(uncompressed))
+      self.evaluate(uncompressed)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations()))
+  def testCompressionVariantMismatch(self):
+    # Use a dataset as a variant.
+    dataset = dataset_ops.Dataset.range(10)
+    variant = dataset._variant_tensor
+    with self.assertRaises(errors.InvalidArgumentError):
+      uncompressed = compression_ops.uncompress(variant, dataset.element_spec)
+      self.evaluate(uncompressed)
+
+  # Only test eager mode since nested datasets are not allowed in graph mode.
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations()))
+  def testDatasetVariantMismatch(self):
+    # Use a nested dataset as an example of a variant.
+    dataset = dataset_ops.Dataset.from_tensors(dataset_ops.Dataset.range(10))
+    with self.assertRaises(TypeError):
+      dataset = dataset.map(
+          lambda x: compression_ops.uncompress(x, dataset.element_spec))
+      self.getDatasetOutput(dataset)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,585 @@
+# 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.copy_to_device()`."""
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.data.experimental.ops import prefetching_ops
+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 iterator_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.util import structure
+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 test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat as util_compat
+
+
+# TODO(b/117581999): add eager coverage when supported.
+class CopyToDeviceTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceHostOptimizations(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    host_dataset = host_dataset.apply(testing.assert_next(["MapAndBatch"]))
+    host_dataset = host_dataset.map(lambda x: x*x).batch(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      self.assertAllEqual([x*x for x in range(10)], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceInt32(self):
+    host_dataset = dataset_ops.Dataset.from_tensors([0, 1, 2, 3])
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int32, next_element.dtype)
+    self.assertEqual((4,), next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      self.assertAllEqual([0, 1, 2, 3], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToSameDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:0"))
+
+    with ops.device("/cpu:0"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceWithPrefetch(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1")).prefetch(1)
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyDictToDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x})
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element["a"].dtype)
+    self.assertEqual([], next_element["a"].shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual({"a": i}, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyDictToDeviceWithPrefetch(self):
+    host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x})
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1")).prefetch(1)
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element["a"].dtype)
+    self.assertEqual([], next_element["a"].shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual({"a": i}, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopySparseTensorsToDevice(self):
+
+    def make_tensor(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0]], values=(i * [1]), dense_shape=[2, 2])
+
+    host_dataset = dataset_ops.Dataset.range(10).map(make_tensor)
+
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        actual = self.evaluate(next_element)
+        self.assertAllEqual([i], actual.values)
+        self.assertAllEqual([[0, 0]], actual.indices)
+        self.assertAllEqual([2, 2], actual.dense_shape)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopySparseTensorsToDeviceWithPrefetch(self):
+
+    def make_tensor(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0]], values=(i * [1]), dense_shape=[2, 2])
+
+    host_dataset = dataset_ops.Dataset.range(10).map(make_tensor)
+
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1")).prefetch(1)
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        actual = self.evaluate(next_element)
+        self.assertAllEqual([i], actual.values)
+        self.assertAllEqual([[0, 0]], actual.indices)
+        self.assertAllEqual([2, 2], actual.dense_shape)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCopyToDeviceGpu(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    with ops.device("/gpu:0"):
+      self.assertDatasetProduces(device_dataset, list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCopyToDeviceGpuWithPrefetch(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0")).prefetch(1)
+
+    with ops.device("/gpu:0"):
+      self.assertDatasetProduces(device_dataset, list(range(10)))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuWithMap(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    def generator():
+      for i in range(10):
+        yield i, float(i), str(i)
+
+    host_dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=(dtypes.int32, dtypes.float32, dtypes.string))
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    def gpu_map_func(x, y, z):
+      return math_ops.square(x), math_ops.square(y), z
+
+    device_dataset = device_dataset.apply(
+        prefetching_ops.map_on_gpu(gpu_map_func))
+    options = options_lib.Options()
+    options.autotune.enabled = False
+    device_dataset = device_dataset.with_options(options)
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        x, y, z = self.evaluate(next_element)
+        self.assertEqual(i**2, x)
+        self.assertEqual(float(i**2), y)
+        self.assertEqual(util_compat.as_bytes(str(i)), z)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuInt32(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.from_tensors([0, 1, 2, 3])
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      self.assertAllEqual([0, 1, 2, 3], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuInt32AndPrefetch(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.from_tensors([0, 1, 2, 3])
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0")).prefetch(1)
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      self.assertAllEqual([0, 1, 2, 3], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuStrings(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.from_tensors(["a", "b", "c"])
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      self.assertAllEqual([b"a", b"b", b"c"], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuStringsAndPrefetch(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.from_tensors(["a", "b", "c"])
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      self.assertAllEqual([b"a", b"b", b"c"], self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDevicePingPongCPUGPU(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0", source_device="/cpu:0"))
+    back_to_cpu_dataset = device_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:0", source_device="/gpu:0"))
+
+    with ops.device("/cpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(back_to_cpu_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceWithReInit(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceWithReInitAndPrefetch(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/cpu:1")).prefetch(1)
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(
+        structure.are_compatible(
+            dataset_ops.get_structure(host_dataset),
+            dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuWithReInit(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testCopyToDeviceGpuWithReInitAndPrefetch(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0")).prefetch(1)
+
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testIteratorGetNextAsOptionalOnGPU(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(3)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.copy_to_device("/gpu:0"))
+    with ops.device("/gpu:0"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_elem = iterator_ops.get_next_as_optional(iterator)
+      elem_has_value_t = next_elem.has_value()
+      elem_value_t = next_elem.get_value()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      # Before initializing the iterator, evaluating the optional fails with
+      # a FailedPreconditionError.
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(elem_has_value_t)
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(elem_value_t)
+
+      # For each element of the dataset, assert that the optional evaluates to
+      # the expected value.
+      self.evaluate(iterator.initializer)
+      for i in range(3):
+        elem_has_value, elem_value = self.evaluate(
+            [elem_has_value_t, elem_value_t])
+        self.assertTrue(elem_has_value)
+        self.assertEqual(i, 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)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,644 @@
+#  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.experimental.CsvDataset`."""
+
+import gzip
+import os
+import zlib
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import error_ops
+from tensorflow.python.data.experimental.ops import readers
+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 as core_readers
+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.ops import parsing_ops
+from tensorflow.python.platform import test
+
+
+class CsvDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _setup_files(self, inputs, linebreak='\n', compression_type=None):
+    filenames = []
+    for i, file_rows in enumerate(inputs):
+      fn = os.path.join(self.get_temp_dir(), 'temp_%d.csv' % i)
+      contents = linebreak.join(file_rows).encode('utf-8')
+      if compression_type is None:
+        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)
+      filenames.append(fn)
+    return filenames
+
+  def _make_test_datasets(self, inputs, **kwargs):
+    # Test by comparing its output to what we could get with map->decode_csv
+    filenames = self._setup_files(inputs)
+    dataset_expected = core_readers.TextLineDataset(filenames)
+    dataset_expected = dataset_expected.map(
+        lambda l: parsing_ops.decode_csv(l, **kwargs))
+    dataset_actual = readers.CsvDataset(filenames, **kwargs)
+    return (dataset_actual, dataset_expected)
+
+  def _test_by_comparison(self, inputs, **kwargs):
+    """Checks that CsvDataset is equiv to TextLineDataset->map(decode_csv)."""
+    dataset_actual, dataset_expected = self._make_test_datasets(
+        inputs, **kwargs)
+    self.assertDatasetsEqual(dataset_actual, dataset_expected)
+
+  def _test_dataset(
+      self,
+      inputs,
+      expected_output=None,
+      expected_err_re=None,
+      linebreak='\n',
+      compression_type=None,  # Used for both setup and parsing
+      **kwargs):
+    """Checks that elements produced by CsvDataset match expected output."""
+    # Convert str type because py3 tf strings are bytestrings
+    filenames = self._setup_files(inputs, linebreak, compression_type)
+    kwargs['compression_type'] = compression_type
+    if expected_err_re is not None:
+      # Verify that OpError is produced as expected
+      with self.assertRaisesOpError(expected_err_re):
+        dataset = readers.CsvDataset(filenames, **kwargs)
+        self.getDatasetOutput(dataset)
+    else:
+      dataset = readers.CsvDataset(filenames, **kwargs)
+      expected_output = [
+          tuple(v.encode('utf-8') if isinstance(v, str) else v
+                for v in op)
+          for op in expected_output
+      ]
+      self.assertDatasetProduces(dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRequiredFields(self):
+    record_defaults = [[]] * 4
+    inputs = [['1,2,3,4']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInt(self):
+    record_defaults = [[0]] * 4
+    inputs = [['1,2,3,4', '5,6,7,8']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFloat(self):
+    record_defaults = [[0.0]] * 4
+    inputs = [['1.0,2.1,3.2,4.3', '5.4,6.5,7.6,8.7']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testString(self):
+    record_defaults = [['']] * 4
+    inputs = [['1.0,2.1,hello,4.3', '5.4,6.5,goodbye,8.7']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithEmptyFields(self):
+    record_defaults = [[0]] * 4
+    inputs = [[',,,', '1,1,1,', ',2,2,2']]
+    self._test_dataset(
+        inputs, [[0, 0, 0, 0], [1, 1, 1, 0], [0, 2, 2, 2]],
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrWithUnquotedQuotes(self):
+    record_defaults = [['']] * 3
+    inputs = [['1,2"3,4']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Unquoted fields cannot have quotes inside',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrWithUnescapedQuotes(self):
+    record_defaults = [['']] * 3
+    inputs = [['"a"b","c","d"']]
+    self._test_dataset(
+        inputs,
+        expected_err_re=
+        'Quote inside a string has to be escaped by another quote',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIgnoreErrWithUnescapedQuotes(self):
+    record_defaults = [['']] * 3
+    inputs = [['1,"2"3",4', '1,"2"3",4",5,5', 'a,b,"c"d"', 'e,f,g']]
+    filenames = self._setup_files(inputs)
+    dataset = readers.CsvDataset(filenames, record_defaults=record_defaults)
+    dataset = dataset.apply(error_ops.ignore_errors())
+    self.assertDatasetProduces(dataset, [(b'e', b'f', b'g')])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIgnoreErrWithUnquotedQuotes(self):
+    record_defaults = [['']] * 3
+    inputs = [['1,2"3,4', 'a,b,c"d', '9,8"7,6,5', 'e,f,g']]
+    filenames = self._setup_files(inputs)
+    dataset = readers.CsvDataset(filenames, record_defaults=record_defaults)
+    dataset = dataset.apply(error_ops.ignore_errors())
+    self.assertDatasetProduces(dataset, [(b'e', b'f', b'g')])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNoQuoteDelimAndUnquotedQuotes(self):
+    record_defaults = [['']] * 3
+    inputs = [['1,2"3,4']]
+    self._test_by_comparison(
+        inputs, record_defaults=record_defaults, use_quote_delim=False)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMixedTypes(self):
+    record_defaults = [
+        constant_op.constant([], dtype=dtypes.int32),
+        constant_op.constant([], dtype=dtypes.float32),
+        constant_op.constant([], dtype=dtypes.string),
+        constant_op.constant([], dtype=dtypes.float64)
+    ]
+    inputs = [['1,2.1,3.2,4.3', '5,6.5,7.6,8.7']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithUseQuoteDelimFalse(self):
+    record_defaults = [['']] * 4
+    inputs = [['1,2,"3,4"', '"5,6",7,8']]
+    self._test_by_comparison(
+        inputs, record_defaults=record_defaults, use_quote_delim=False)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithFieldDelim(self):
+    record_defaults = [[0]] * 4
+    inputs = [['1:2:3:4', '5:6:7:8']]
+    self._test_by_comparison(
+        inputs, record_defaults=record_defaults, field_delim=':')
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNaValue(self):
+    record_defaults = [[0]] * 4
+    inputs = [['1,NA,3,4', 'NA,6,7,8']]
+    self._test_by_comparison(
+        inputs, record_defaults=record_defaults, na_value='NA')
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithSelectCols(self):
+    record_defaults = [['']] * 2
+    inputs = [['1,2,3,4', '"5","6","7","8"']]
+    self._test_by_comparison(
+        inputs, record_defaults=record_defaults, select_cols=[1, 2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithSelectColsTooHigh(self):
+    record_defaults = [[0]] * 2
+    inputs = [['1,2,3,4', '5,6,7,8']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 2 fields but have 1 in record',
+        record_defaults=record_defaults,
+        select_cols=[3, 4])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithOneCol(self):
+    record_defaults = [['NA']]
+    inputs = [['0', '', '2']]
+    self._test_dataset(
+        inputs, [['0'], ['NA'], ['2']], record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithMultipleFiles(self):
+    record_defaults = [[0]] * 4
+    inputs = [['1,2,3,4', '5,6,7,8'], ['5,6,7,8']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithLeadingAndTrailingSpaces(self):
+    record_defaults = [[0.0]] * 4
+    inputs = [['0, 1, 2, 3']]
+    expected = [[0.0, 1.0, 2.0, 3.0]]
+    self._test_dataset(inputs, expected, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithMissingDefault(self):
+    record_defaults = [[]] * 2
+    inputs = [['0,']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Field 1 is required but missing in record!',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithFewerDefaultsThanFields(self):
+    record_defaults = [[0.0]] * 2
+    inputs = [['0,1,2,3']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 2 fields but have more in record',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithMoreDefaultsThanFields(self):
+    record_defaults = [[0.0]] * 5
+    inputs = [['0,1,2,3']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 5 fields but have 4 in record',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithHeader(self):
+    record_defaults = [[0]] * 2
+    inputs = [['col1,col2', '1,2']]
+    expected = [[1, 2]]
+    self._test_dataset(
+        inputs,
+        expected,
+        record_defaults=record_defaults,
+        header=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithHeaderAndNoRecords(self):
+    record_defaults = [[0]] * 2
+    inputs = [['col1,col2']]
+    expected = []
+    self._test_dataset(
+        inputs,
+        expected,
+        record_defaults=record_defaults,
+        header=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithHeaderEmptyFile(self):
+    record_defaults = [[0]] * 2
+    inputs = [[]]
+    expected_err_re = "Can't read header of file"
+    self._test_dataset(
+        inputs,
+        expected_err_re=expected_err_re,
+        record_defaults=record_defaults,
+        header=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithEmptyFile(self):
+    record_defaults = [['']] * 2
+    inputs = [['']]  # Empty file
+    self._test_dataset(
+        inputs, expected_output=[], record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithEmptyRecord(self):
+    record_defaults = [['']] * 2
+    inputs = [['', '1,2']]  # First record is empty
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 2 fields but have 1 in record',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithChainedOps(self):
+    # Testing that one dataset can create multiple iterators fine.
+    # `repeat` creates multiple iterators from the same C++ Dataset.
+    record_defaults = [[0]] * 4
+    inputs = [['1,,3,4', '5,6,,8']]
+    ds_actual, ds_expected = self._make_test_datasets(
+        inputs, record_defaults=record_defaults)
+    self.assertDatasetsEqual(
+        ds_actual.repeat(5).prefetch(1),
+        ds_expected.repeat(5).prefetch(1))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithTypeDefaults(self):
+    # Testing using dtypes as record_defaults for required fields
+    record_defaults = [dtypes.float32, [0.0]]
+    inputs = [['1.0,2.0', '3.0,4.0']]
+    self._test_dataset(
+        inputs,
+        [[1.0, 2.0], [3.0, 4.0]],
+        record_defaults=record_defaults,
+    )
+
+## The following tests exercise parsing logic for quoted fields
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithQuoted(self):
+    record_defaults = [['']] * 4
+    inputs = [['"a","b","c :)","d"', '"e","f","g :(","h"']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithOneColAndQuotes(self):
+    record_defaults = [['']]
+    inputs = [['"0"', '"1"', '"2"']]
+    self._test_dataset(
+        inputs, [['0'], ['1'], ['2']], record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNewLine(self):
+    # In this case, we expect it to behave differently from
+    # TextLineDataset->map(decode_csv) since that flow has bugs
+    record_defaults = [['']] * 4
+    inputs = [['a,b,"""c""\n0","d\ne"', 'f,g,h,i']]
+    expected = [['a', 'b', '"c"\n0', 'd\ne'], ['f', 'g', 'h', 'i']]
+    self._test_dataset(inputs, expected, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNewLineInUnselectedCol(self):
+    record_defaults = [['']]
+    inputs = [['1,"2\n3",4', '5,6,7']]
+    self._test_dataset(
+        inputs,
+        expected_output=[['1'], ['5']],
+        record_defaults=record_defaults,
+        select_cols=[0])
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testWithExcludeCol(self):
+    record_defaults = [['']]
+    inputs = [['1,2,3', '5,6,7']]
+    self._test_dataset(
+        inputs,
+        expected_output=[['1'], ['5']],
+        record_defaults=record_defaults,
+        exclude_cols=[1, 2])
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testWithSelectandExcludeCol(self):
+    record_defaults = [['']]
+    inputs = [['1,2,3', '5,6,7']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Either select_cols or exclude_cols should be empty',
+        record_defaults=record_defaults,
+        select_cols=[0],
+        exclude_cols=[1, 2])
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testWithExcludeColandRecordDefaultsTooLow(self):
+    record_defaults = [['']]
+    inputs = [['1,2,3', '5,6,7']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 1 fields but have more in record',
+        record_defaults=record_defaults,
+        exclude_cols=[0])
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testWithExcludeColandRecordDefaultsTooHigh(self):
+    record_defaults = [['']] * 3
+    inputs = [['1,2,3', '5,6,7']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='Expect 3 fields but have 2 in record',
+        record_defaults=record_defaults,
+        exclude_cols=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithMultipleNewLines(self):
+    # In this case, we expect it to behave differently from
+    # TextLineDataset->map(decode_csv) since that flow has bugs
+    record_defaults = [['']] * 4
+    inputs = [['a,"b\n\nx","""c""\n \n0","d\ne"', 'f,g,h,i']]
+    expected = [['a', 'b\n\nx', '"c"\n \n0', 'd\ne'], ['f', 'g', 'h', 'i']]
+    self._test_dataset(inputs, expected, record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorWithTerminateMidRecord(self):
+    record_defaults = [['']] * 4
+    inputs = [['a,b,c,"a']]
+    self._test_dataset(
+        inputs,
+        expected_err_re=
+        'Reached end of file without closing quoted field in record',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithEscapedQuotes(self):
+    record_defaults = [['']] * 4
+    inputs = [['1.0,2.1,"she said: ""hello""",4.3', '5.4,6.5,goodbye,8.7']]
+    self._test_by_comparison(inputs, record_defaults=record_defaults)
+
+
+## Testing that parsing works with all buffer sizes, quoted/unquoted fields,
+## and different types of line breaks
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithInvalidBufferSize(self):
+    record_defaults = [['']] * 4
+    inputs = [['a,b,c,d']]
+    self._test_dataset(
+        inputs,
+        expected_err_re='buffer_size should be positive',
+        record_defaults=record_defaults,
+        buffer_size=0)
+
+  def _test_dataset_on_buffer_sizes(self,
+                                    inputs,
+                                    expected,
+                                    linebreak,
+                                    record_defaults,
+                                    compression_type=None,
+                                    num_sizes_to_test=20):
+    # Testing reading with a range of buffer sizes that should all work.
+    for i in list(range(1, 1 + num_sizes_to_test)) + [None]:
+      self._test_dataset(
+          inputs,
+          expected,
+          linebreak=linebreak,
+          compression_type=compression_type,
+          record_defaults=record_defaults,
+          buffer_size=i)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithLF(self):
+    record_defaults = [['NA']] * 3
+    inputs = [['abc,def,ghi', '0,1,2', ',,']]
+    expected = [['abc', 'def', 'ghi'], ['0', '1', '2'], ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\n', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCR(self):
+    # Test that when the line separator is '\r', parsing works with all buffer
+    # sizes
+    record_defaults = [['NA']] * 3
+    inputs = [['abc,def,ghi', '0,1,2', ',,']]
+    expected = [['abc', 'def', 'ghi'], ['0', '1', '2'], ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\r', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCRLF(self):
+    # Test that when the line separator is '\r\n', parsing works with all buffer
+    # sizes
+    record_defaults = [['NA']] * 3
+    inputs = [['abc,def,ghi', '0,1,2', ',,']]
+    expected = [['abc', 'def', 'ghi'], ['0', '1', '2'], ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\r\n', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithBufferSizeAndQuoted(self):
+    record_defaults = [['NA']] * 3
+    inputs = [['"\n\n\n","\r\r\r","abc"', '"0","1","2"', '"","",""']]
+    expected = [['\n\n\n', '\r\r\r', 'abc'], ['0', '1', '2'],
+                ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\n', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCRAndQuoted(self):
+    # Test that when the line separator is '\r', parsing works with all buffer
+    # sizes
+    record_defaults = [['NA']] * 3
+    inputs = [['"\n\n\n","\r\r\r","abc"', '"0","1","2"', '"","",""']]
+    expected = [['\n\n\n', '\r\r\r', 'abc'], ['0', '1', '2'],
+                ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\r', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCRLFAndQuoted(self):
+    # Test that when the line separator is '\r\n', parsing works with all buffer
+    # sizes
+    record_defaults = [['NA']] * 3
+    inputs = [['"\n\n\n","\r\r\r","abc"', '"0","1","2"', '"","",""']]
+    expected = [['\n\n\n', '\r\r\r', 'abc'], ['0', '1', '2'],
+                ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs, expected, linebreak='\r\n', record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithGzipCompressionType(self):
+    record_defaults = [['NA']] * 3
+    inputs = [['"\n\n\n","\r\r\r","abc"', '"0","1","2"', '"","",""']]
+    expected = [['\n\n\n', '\r\r\r', 'abc'], ['0', '1', '2'],
+                ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs,
+        expected,
+        linebreak='\r\n',
+        compression_type='GZIP',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithZlibCompressionType(self):
+    record_defaults = [['NA']] * 3
+    inputs = [['"\n\n\n","\r\r\r","abc"', '"0","1","2"', '"","",""']]
+    expected = [['\n\n\n', '\r\r\r', 'abc'], ['0', '1', '2'],
+                ['NA', 'NA', 'NA']]
+    self._test_dataset_on_buffer_sizes(
+        inputs,
+        expected,
+        linebreak='\r\n',
+        compression_type='ZLIB',
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithScalarDefaults(self):
+    record_defaults = [constant_op.constant(0, dtype=dtypes.int64)] * 4
+    inputs = [[',,,', '1,1,1,', ',2,2,2']]
+    self._test_dataset(
+        inputs, [[0, 0, 0, 0], [1, 1, 1, 0], [0, 2, 2, 2]],
+        record_defaults=record_defaults)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWith2DDefaults(self):
+    record_defaults = [constant_op.constant([[0]], dtype=dtypes.int64)] * 4
+    inputs = [[',,,', '1,1,1,', ',2,2,2']]
+
+    if context.executing_eagerly():
+      err_spec = errors.InvalidArgumentError, (
+          'Each record default should be at '
+          'most rank 1')
+    else:
+      err_spec = ValueError, 'Shape must be at most rank 1 but is rank 2'
+
+    with self.assertRaisesWithPredicateMatch(*err_spec):
+      self._test_dataset(
+          inputs, [[0, 0, 0, 0], [1, 1, 1, 0], [0, 2, 2, 2]],
+          record_defaults=record_defaults)
+
+  def testImmutableParams(self):
+    inputs = [['a,b,c', '1,2,3', '4,5,6']]
+    filenames = self._setup_files(inputs)
+    select_cols = ['a', 'c']
+    _ = readers.make_csv_dataset(
+        filenames, batch_size=1, select_columns=select_cols)
+    self.assertAllEqual(select_cols, ['a', 'c'])
+
+
+class CsvDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                               parameterized.TestCase):
+
+  def setUp(self):
+    super(CsvDatasetCheckpointTest, self).setUp()
+    self._num_cols = 7
+    self._num_rows = 10
+    self._num_epochs = 14
+    self._num_outputs = self._num_rows * self._num_epochs
+
+    inputs = [
+        ','.join(str(self._num_cols * j + i)
+                 for i in range(self._num_cols))
+        for j in range(self._num_rows)
+    ]
+    contents = '\n'.join(inputs).encode('utf-8')
+
+    self._filename = os.path.join(self.get_temp_dir(), 'file.csv')
+    self._compressed = os.path.join(self.get_temp_dir(),
+                                    'comp.csv')  # GZip compressed
+
+    with open(self._filename, 'wb') as f:
+      f.write(contents)
+    with gzip.GzipFile(self._compressed, 'wb') as f:
+      f.write(contents)
+
+  def ds_func(self, **kwargs):
+    compression_type = kwargs.get('compression_type', None)
+    if compression_type == 'GZIP':
+      filename = self._compressed
+    elif compression_type is None:
+      filename = self._filename
+    else:
+      raise ValueError('Invalid compression type:', compression_type)
+
+    return readers.CsvDataset(filename, **kwargs).repeat(self._num_epochs)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testCore(self, verify_fn):
+    defs = [[0]] * self._num_cols
+    verify_fn(self, lambda: self.ds_func(record_defaults=defs, buffer_size=2),
+              self._num_outputs)
+
+
+if __name__ == '__main__':
+  test.main()

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

@@ -0,0 +1,128 @@
+# 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.dense_to_sparse_batch()."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import batching
+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)).apply(
+            batching.dense_to_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)).apply(
+            batching.dense_to_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).apply(
+          batching.dense_to_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).apply(
+          batching.dense_to_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)).apply(
+            batching.dense_to_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/experimental/kernel_tests/from_list_test.py

@@ -0,0 +1,303 @@
+# Copyright 2022 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.from_list()."""
+
+import collections
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import from_list
+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 tensor_shape
+from tensorflow.python.platform import test
+
+
+class FromListTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(elements=[
+              combinations.NamedObject("empty_input", []),
+              combinations.NamedObject("non-list_input", (1, 2, 3)),
+          ])))
+  def testInvalidInputs(self, elements):
+    with self.assertRaises(ValueError):
+      from_list.from_list(elements._obj)
+
+  def testLargeNInputs(self):
+    elements = list(range(1000))
+    dataset = from_list.from_list(elements)
+    self.assertDatasetProduces(dataset, expected_output=elements)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTupleInputs(self):
+    elements = [(1, 2), (3, 4)]
+    dataset = from_list.from_list(elements)
+    self.assertEqual(
+        [np.shape(c) for c in elements[0]],
+        [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    self.assertDatasetProduces(dataset, expected_output=elements)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(map_fn=[
+              combinations.NamedObject("simple", lambda x: x),
+              combinations.NamedObject(
+                  "ordered_dicts",
+                  lambda x: collections.OrderedDict([("x", x)]))
+          ])))
+  def testDatasetInputs(self, map_fn):
+    dss = [dataset_ops.Dataset.range(10).map(map_fn) for _ in range(10)]
+    ds = dataset_ops.Dataset.from_tensor_slices(dss)
+    ds = ds.flat_map(lambda x: x)
+    self.assertDatasetProduces(
+        ds, expected_output=[map_fn(x) for x in list(range(10)) * 10])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNonRectangularInputs(self):
+    elements = [[[1]], [[2, 3]], [[4, 5, 6]]]
+    dataset = from_list.from_list(elements)
+    self.assertEqual(
+        tensor_shape.Dimension(1),
+        dataset_ops.get_legacy_output_shapes(dataset)[0])
+    self.assertDatasetProduces(dataset, expected_output=elements)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDictInputs(self):
+    elements = [{
+        "foo": [1, 2, 3],
+        "bar": [[4.0], [5.0], [6.0]]
+    }, {
+        "foo": [4, 5, 6],
+        "bar": [[7.0], [8.0], [9.0]]
+    }]
+    dataset = from_list.from_list(elements)
+    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((3,), dataset_ops.get_legacy_output_shapes(dataset)["foo"])
+    self.assertEqual((3, 1),
+                     dataset_ops.get_legacy_output_shapes(dataset)["bar"])
+    self.assertDatasetProduces(dataset, expected_output=elements)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUintInputs(self):
+    elements = [(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))]
+    dataset = from_list.from_list(elements)
+    self.assertEqual(
+        (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64),
+        dataset_ops.get_legacy_output_types(dataset))
+    self.assertDatasetProduces(dataset, elements)
+
+
+class FromListRandomAccessTest(test_base.DatasetTestBase,
+                               parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidIndex(self):
+    dataset = from_list.from_list([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 = from_list.from_list(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 = from_list.from_list(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 testMultipleElements(self):
+    dataset = from_list.from_list([[1, 2], [3, 4], [5, 6]])
+    self.assertEqual(1, self.evaluate(random_access.at(dataset, 0))[0])
+    self.assertEqual(2, self.evaluate(random_access.at(dataset, 0))[1])
+    self.assertEqual(3, self.evaluate(random_access.at(dataset, 1))[0])
+    self.assertEqual(4, self.evaluate(random_access.at(dataset, 1))[1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDictionary(self):
+    dataset = from_list.from_list([{"a": 1, "b": 3}, {"a": 2, "b": 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):
+    elements = [
+        np.tile(np.array([[0], [1]], dtype=np.uint64), 2),
+        np.tile(np.array([[2], [256]], dtype=np.uint64), 2),
+        np.tile(np.array([[4], [65536]], dtype=np.uint64), 2),
+        np.tile(np.array([[8], [4294967296]], dtype=np.uint64), 2),
+    ]
+    dataset = from_list.from_list(elements)
+    for i in range(len(elements)):
+      result = self.evaluate(random_access.at(dataset, i))
+      self.assertAllEqual(elements[i], result)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testName(self):
+    dataset = from_list.from_list([42], name="from_list")
+    self.assertDatasetProduces(dataset, [42])
+
+
+class FromListCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                             parameterized.TestCase):
+
+  def _build_list_dataset(self, elements, options=None):
+    dataset = from_list.from_list(elements)
+    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 elements
+    elements = [
+        np.tile(np.array([[1], [2], [3], [4]]), 20),
+        np.tile(np.array([[12], [13], [14], [15]]), 22),
+        np.array([37, 38, 39, 40])
+    ]
+    options = options_lib.Options()
+    options.experimental_symbolic_checkpoint = symbolic_checkpoint
+    verify_fn(
+        self,
+        lambda: self._build_list_dataset(elements, options),
+        num_outputs=3)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testDict(self, verify_fn):
+    dict_elements = [{
+        "foo": 1,
+        "bar": 4.0
+    }, {
+        "foo": 2,
+        "bar": 5.0
+    }, {
+        "foo": 3,
+        "bar": 6.0
+    }]
+    verify_fn(
+        self, lambda: self._build_list_dataset(dict_elements), num_outputs=3)
+
+
+class FromListGlobalShuffleTest(
+    test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[100],
+              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 = from_list.from_list(list(range(dataset_range)))
+    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(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 FromListGlobalShuffleCheckpointTest(
+    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],
+              symbolic_checkpoint=[True, False])))
+  def test(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      repetitions: int,
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = from_list.from_list(list(range(dataset_range)))
+      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 * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,326 @@
+# 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 global shuffling of tf.data datasets."""
+
+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.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 random_seed
+from tensorflow.python.platform import test
+
+
+class GlobalShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for global shuffling of tf.data datasets."""
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              dataset_range=[1, 100],
+              seed=[None, 42],
+              use_tensor_seed=[True, False],
+              prefetch=[True, False])))
+  def testRange(
+      self,
+      dataset_range: int,
+      seed: Optional[int],
+      use_tensor_seed: bool,
+      prefetch: bool):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    if prefetch:
+      dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    seed = (constant_op.constant(seed, dtype=dtypes.int64)
+            if seed and use_tensor_seed else seed)
+    dataset = global_shuffle_op._global_shuffle(dataset, seed=seed)
+    dataset = dataset.repeat(3)
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(output, list(range(dataset_range)) * 3)
+
+    output_per_iteration = [
+        output[i : i + dataset_range]
+        for i in range(0, len(output), dataset_range)]
+    self.assertCountEqual(output_per_iteration[0], list(range(dataset_range)))
+    self.assertCountEqual(output_per_iteration[1], list(range(dataset_range)))
+    self.assertCountEqual(output_per_iteration[2], list(range(dataset_range)))
+    if dataset_range > 1:
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[1])
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[2])
+      self.assertNotEqual(output_per_iteration[1], output_per_iteration[2])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(dataset_range=[1, 100], seed=[None, 42])))
+  def testNegativeRange(self, dataset_range: int, seed: Optional[int]):
+    dataset = dataset_ops.Dataset.range(dataset_range, -dataset_range, -1)
+    dataset = global_shuffle_op._global_shuffle(dataset)
+    dataset = dataset.repeat(3)
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(
+        output, list(range(dataset_range, -dataset_range, -1)) * 3)
+
+    output_per_iteration = [
+        output[i : i + dataset_range * 2]
+        for i in range(0, len(output), dataset_range * 2)]
+    self.assertCountEqual(output_per_iteration[0],
+                          list(range(dataset_range, -dataset_range, -1)))
+    self.assertCountEqual(output_per_iteration[1],
+                          list(range(dataset_range, -dataset_range, -1)))
+    self.assertCountEqual(output_per_iteration[2],
+                          list(range(dataset_range, -dataset_range, -1)))
+    if dataset_range > 1:
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[1])
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[2])
+      self.assertNotEqual(output_per_iteration[1], output_per_iteration[2])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(reshuffle=[True, False], seed=[None, 42])))
+  def testReshuffleRepeatEpochs(self, reshuffle: bool, seed: Optional[int]):
+    dataset_range = 100
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle)
+    dataset = dataset.repeat(2)
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(output, list(range(dataset_range)) * 2)
+    output_per_iteration = [
+        output[i : i + dataset_range]
+        for i in range(0, len(output), dataset_range)]
+    if reshuffle:
+      self.assertNotEqual(output_per_iteration[0], output_per_iteration[1])
+    else:
+      self.assertEqual(output_per_iteration[0], output_per_iteration[1])
+
+  # 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"),
+          combinations.combine(reshuffle=[True, False], seed=[None, 42])))
+  def testReshuffleIterationEpochs(self, reshuffle: bool, seed: Optional[int]):
+    # 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_range = 100
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = global_shuffle_op._global_shuffle(
+        dataset, seed=seed, reshuffle_each_iteration=reshuffle)
+
+    first_epoch = self.getDatasetOutput(dataset)
+    second_epoch = self.getDatasetOutput(dataset)
+    if reshuffle:
+      self.assertNotEqual(first_epoch, second_epoch)
+    else:
+      self.assertEqual(first_epoch, second_epoch)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    dataset = dataset_ops.Dataset.range(0)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "`global_shuffle` requires the input dataset to have a non-empty "
+        "finite cardinality."):
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnsupportedDataset(self):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset.shuffle(buffer_size=1)
+    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 GlobalShuffleCheckpointTest(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=[1, 10],
+              reshuffle_each_iteration=[True, False],
+              prefetch=[True, False],
+              symbolic_checkpoint=[True, False])))
+  def testRange(
+      self,
+      verify_fn: Callable[..., None],
+      dataset_range: int,
+      reshuffle_each_iteration: bool,
+      prefetch: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      if prefetch:
+        dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      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,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+class GlobalShuffleNumpyIteratorCheckpointTest(parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.v2_eager_only_combinations(),
+          combinations.combine(
+              dataset_range=[6, 10],
+              reshuffle_each_iteration=[True, False],
+              prefetch=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testRange(
+      self,
+      dataset_range: int,
+      reshuffle_each_iteration: bool,
+      prefetch: bool,
+      symbolic_checkpoint: bool,
+  ):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      if prefetch:
+        dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=10, 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
+
+    dataset = _build_dataset()
+    it = dataset.as_numpy_iterator()
+
+    expected_elements = []
+    # For each step, we checkpoint the iterator to try to regenerate
+    # the element again to make sure the regenerated element
+    # is the same as expected.
+    for _ in range(dataset_range):
+      checkpoint = it.save()
+      expected = next(it)
+      it.restore(checkpoint)
+      actual = next(it)
+
+      self.assertEqual(expected, actual)
+
+      expected_elements.append(expected)
+
+    self.assertCountEqual(expected_elements, range(dataset_range))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.v2_eager_only_combinations(),
+          combinations.combine(
+              dataset_range=[6, 10],
+              reshuffle_each_iteration=[True, False],
+              prefetch=[True, False],
+              symbolic_checkpoint=[True, False],
+          ),
+      )
+  )
+  def testRangeAndRepeat(
+      self,
+      dataset_range: int,
+      reshuffle_each_iteration: bool,
+      prefetch: bool,
+      symbolic_checkpoint: bool,
+  ):
+
+    num_repeat = 3
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.range(dataset_range)
+      if prefetch:
+        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.repeat(num_repeat)
+      if symbolic_checkpoint:
+        options = options_lib.Options()
+        options.experimental_symbolic_checkpoint = symbolic_checkpoint
+        dataset = dataset.with_options(options)
+      return dataset
+
+    dataset = _build_dataset()
+    it = dataset.as_numpy_iterator()
+    checkpoint = it.save()
+    first_pass_results = []
+
+    # First pass
+    for _ in range(dataset_range * num_repeat):
+      e = next(it)
+      first_pass_results.append(e)
+
+    it.restore(checkpoint)
+
+    # Second pass
+    second_pass_results = []
+    for _ in range(dataset_range * num_repeat):
+      checkpoint = it.save()
+      expected = next(it)
+      it.restore(checkpoint)
+      actual = next(it)
+
+      self.assertEqual(expected, actual)
+
+      second_pass_results.append(expected)
+
+    self.assertCountEqual(
+        second_pass_results, list(range(dataset_range)) * num_repeat
+    )
+
+    self.assertSequenceEqual(
+        first_pass_results,
+        second_pass_results,
+        "First pass and second pass should generate the same results because"
+        " the underlying seed generator should be restored properly as well.",
+    )
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,220 @@
+# 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_reducer()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import grouping
+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 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.platform import test
+
+
+class GroupByReducerTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSum(self):
+    reducer = grouping.Reducer(
+        init_func=lambda _: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).apply(
+          grouping.group_by_reducer(lambda x: x % 2, reducer))
+      self.assertDatasetProduces(
+          dataset,
+          expected_shapes=tensor_shape.TensorShape([]),
+          expected_output=[(i - 1) * i, i * i])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAverage(self):
+
+    def reduce_fn(x, y):
+      return (x[0] * x[1] + math_ops.cast(y, dtypes.float32)) / (
+          x[1] + 1), x[1] + 1
+
+    reducer = grouping.Reducer(
+        init_func=lambda _: (0.0, 0.0),
+        reduce_func=reduce_fn,
+        finalize_func=lambda x, _: x)
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).apply(
+          grouping.group_by_reducer(
+              lambda x: math_ops.cast(x, dtypes.int64) % 2, reducer))
+      self.assertDatasetProduces(
+          dataset,
+          expected_shapes=tensor_shape.TensorShape([]),
+          expected_output=[i - 1, i])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcat(self):
+    components = np.array(list("abcdefghijklmnopqrst")).view(np.char.chararray)
+    reducer = grouping.Reducer(
+        init_func=lambda x: "",
+        reduce_func=lambda x, y: x + y[0],
+        finalize_func=lambda x: x)
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.from_tensor_slices(components),
+           dataset_ops.Dataset.range(2 * i))).apply(
+               grouping.group_by_reducer(lambda x, y: y % 2, reducer))
+      self.assertDatasetProduces(
+          dataset,
+          expected_shapes=tensor_shape.TensorShape([]),
+          expected_output=[b"acegikmoqs"[:i], b"bdfhjlnprt"[:i]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparseSum(self):
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1], dtype=np.int64)),
+          dense_shape=np.array([1, 1]))
+
+    reducer = grouping.Reducer(
+        init_func=lambda _: _sparse(np.int64(0)),
+        reduce_func=lambda x, y: _sparse(x.values[0] + y.values[0]),
+        finalize_func=lambda x: x.values[0])
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).map(_sparse).apply(
+          grouping.group_by_reducer(lambda x: x.values[0] % 2, reducer))
+      self.assertDatasetProduces(
+          dataset,
+          expected_shapes=tensor_shape.TensorShape([]),
+          expected_output=[(i - 1) * i, i * i])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChangingStateShape(self):
+
+    def reduce_fn(x, _):
+      # Statically known rank, but dynamic length.
+      larger_dim = array_ops.concat([x[0], x[0]], 0)
+      # Statically unknown rank.
+      larger_rank = array_ops.expand_dims(x[1], 0)
+      return larger_dim, larger_rank
+
+    reducer = grouping.Reducer(
+        init_func=lambda x: ([0], 1),
+        reduce_func=reduce_fn,
+        finalize_func=lambda x, y: (x, y))
+
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.from_tensors(np.int64(0)).repeat(i).apply(
+          grouping.group_by_reducer(lambda x: x, reducer))
+      dataset_output_shapes = dataset_ops.get_legacy_output_shapes(dataset)
+      self.assertEqual([None], dataset_output_shapes[0].as_list())
+      self.assertIs(None, dataset_output_shapes[1].ndims)
+      get_next = self.getNext(dataset)
+      x, y = self.evaluate(get_next())
+      self.assertAllEqual([0] * (2**i), x)
+      self.assertAllEqual(np.array(1, ndmin=i), y)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTypeMismatch(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: constant_op.constant(1, dtype=dtypes.int32),
+        reduce_func=lambda x, y: constant_op.constant(1, dtype=dtypes.int64),
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaises(TypeError):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: np.int64(0), reducer))
+
+  # TODO(b/78665031): Remove once non-scalar keys are supported.
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidKeyShape(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaises(ValueError):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: np.int64((0, 0)), reducer))
+
+  # TODO(b/78665031): Remove once non-int64 keys are supported.
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidKeyType(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaises(ValueError):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: "wrong", reducer))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTuple(self):
+    def init_fn(_):
+      return np.array([], dtype=np.int64), np.int64(0)
+
+    def reduce_fn(state, value):
+      s1, s2 = state
+      v1, v2 = value
+      return array_ops.concat([s1, [v1]], 0), s2 + v2
+
+    def finalize_fn(s1, s2):
+      return s1, s2
+
+    reducer = grouping.Reducer(init_fn, reduce_fn, finalize_fn)
+    dataset = dataset_ops.Dataset.zip(
+        (dataset_ops.Dataset.range(10), dataset_ops.Dataset.range(10))).apply(
+            grouping.group_by_reducer(lambda x, y: np.int64(0), reducer))
+    get_next = self.getNext(dataset)
+    x, y = self.evaluate(get_next())
+    self.assertAllEqual(x, np.asarray([x for x in range(10)]))
+    self.assertEqual(y, 45)
+
+
+class GroupByReducerCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                   parameterized.TestCase):
+
+  def _build_dataset(self, components):
+    reducer = grouping.Reducer(
+        init_func=lambda _: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    return dataset_ops.Dataset.from_tensor_slices(components).apply(
+        grouping.group_by_reducer(lambda x: x % 5, reducer))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    components = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int64)
+    verify_fn(
+        self,
+        lambda: self._build_dataset(components),
+        num_outputs=5)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,302 @@
+# 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 the index flat map dataset."""
+
+from typing import Any, Callable, Optional, Union
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import global_shuffle_op
+from tensorflow.python.data.experimental.ops import index_flat_map_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 tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.platform import test
+
+_IndexType = index_flat_map_op._IndexType
+
+
+class IndexFlatMapTest(test_base.DatasetTestBase, parameterized.TestCase):
+  """Tests for global shuffling of index flat map datasets."""
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(use_tensors=[True, False])))
+  def test_split_strings(self, use_tensors: bool):
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    # The metadata is [(0, 2, 0), (2, 6, 1), (6, 8, 2), (8, 9, 3)].
+    metadata = _get_metadata(input_data)
+
+    def _index_map_func(index: _IndexType) -> tuple[_IndexType, _IndexType]:
+      index = _maybe_convert_to_tensor(index)
+      element_index, offset = _get_index_map_func(metadata)(index)
+      return (_maybe_convert_to_tensor(element_index),
+              _maybe_convert_to_tensor(offset))
+
+    def _maybe_convert_to_tensor(value: Any) -> _IndexType:
+      return math_ops.cast(value, dtypes.int64) if use_tensors else value
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = index_flat_map_op.index_flat_map(dataset, _split, _index_map_func)
+    output = self.getDatasetOutput(dataset)
+    self.assertEqual(output,
+                     [b"0", b"1", b"2", b"3", b"4", b"5", b"6", b"7", b"8"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_cache(self):
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    metadata = _get_metadata(input_data)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = dataset.cache()
+    dataset = index_flat_map_op.index_flat_map(
+        dataset, _split, _get_index_map_func(metadata))
+    output = self.getDatasetOutput(dataset)
+    self.assertEqual(output,
+                     [b"0", b"1", b"2", b"3", b"4", b"5", b"6", b"7", b"8"])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              repetitions=[1, 3],
+              seed=[None, 42],
+              reshuffle_each_iteration=[True, False])))
+  def test_global_shuffle(
+      self,
+      repetitions: int,
+      seed: Optional[int],
+      reshuffle_each_iteration: bool):
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    metadata = _get_metadata(input_data)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = index_flat_map_op.index_flat_map(
+        dataset, _split, _get_index_map_func(metadata), output_cardinality=9)
+    self.assertEqual(self.evaluate(dataset.cardinality()), 9)
+
+    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 = [
+        b"0", b"1", b"2", b"3", b"4", b"5", b"6", b"7", b"8"] * repetitions
+    self.assertCountEqual(dataset_output, expected)
+    self.assertNotEqual(dataset_output, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(dataset_range=[0, 10])))
+  def test_identity_map(self, dataset_range: int):
+
+    def _map_func(element: Any) -> Any:
+      return element
+
+    def _index_map_func(index: int) -> tuple[int, int]:
+      return (index, 0)
+
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    dataset = index_flat_map_op.index_flat_map(
+        dataset, _map_func, _index_map_func)
+    self.assertDatasetProduces(dataset, list(range(dataset_range)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              input_range=[0, 10],
+              map_output=[
+                  [[1, 2], [3, 4], [5, 6]],
+                  [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]]
+              ],
+              use_tensors=[True, False])))
+  def test_nested_list(
+      self, input_range: int, map_output: list[Any], use_tensors: bool):
+
+    def _map_func(_) -> Union[tensor.Tensor, list[list[int]]]:
+      return (constant_op.constant(map_output, dtype=dtypes.int64)
+              if use_tensors else map_output)
+
+    def _index_map_func(i: int) -> tuple[int, int]:
+      return (i // len(map_output), i % len(map_output))
+
+    dataset = dataset_ops.Dataset.range(input_range)
+    dataset = index_flat_map_op.index_flat_map(
+        dataset, _map_func, _index_map_func)
+    self.assertDatasetProduces(dataset, map_output * input_range)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_offset_out_of_range(self):
+
+    def _index_map_func(_) -> tuple[int, int]:
+      return (0, 1000)
+
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = index_flat_map_op.index_flat_map(dataset, _split, _index_map_func)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "invalid `index_map_func` which returns offset 1000"):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_invalid_map_fn_type(self):
+
+    def _index_map_func(_) -> str:
+      # Expected to return two integers.
+      return "Hello"
+
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = index_flat_map_op.index_flat_map(dataset, _split, _index_map_func)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "expected to return two int values"):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_unknown_cardinality(self):
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "`global_shuffle` requires the input dataset to have a non-empty "
+        "finite cardinality."):
+      dataset = index_flat_map_op.index_flat_map(
+          dataset, _split, _get_index_map_func(_get_metadata(input_data)))
+      dataset = global_shuffle_op._global_shuffle(dataset)
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+class IndexFlatMapCheckpointTest(
+    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],
+              symbolic_checkpoint=[True, False])))
+  def test_index_flat_map(
+      self,
+      verify_fn: Callable[..., None],
+      repetitions: int,
+      symbolic_checkpoint: bool):
+
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+      dataset = index_flat_map_op.index_flat_map(
+          dataset, _split, _get_index_map_func(_get_metadata(input_data)))
+      if repetitions > 1:
+        dataset = dataset.repeat(repetitions)
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(self, _build_dataset, num_outputs=9 * repetitions)
+
+  @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 test_global_shuffle(
+      self,
+      verify_fn: Callable[..., None],
+      repetitions: list[int],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    input_data = ["0 1", "2 3 4 5", "6 7", "8"]
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+      dataset = index_flat_map_op.index_flat_map(
+          dataset,
+          _split,
+          _get_index_map_func(_get_metadata(input_data)),
+          output_cardinality=9)
+      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=9 * repetitions,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+def _split(element: str) -> tensor.Tensor:
+  return ragged_string_ops.string_split_v2(element, " ")
+
+
+def _get_metadata(input_data: list[str]) -> tensor.Tensor:
+  """Given a list of strings, creates a metadata matrix."""
+
+  metadata = []
+  for i, data in enumerate(input_data):
+    split_data = data.split()
+    last_index = metadata[-1][1] if metadata else 0
+    metadata.append((last_index, last_index + len(split_data), i))
+  return constant_op.constant(metadata, dtype=dtypes.int64)
+
+
+def _get_index_map_func(
+    metadata: tensor.Tensor) -> Callable[[int], tuple[int, int]]:
+  """Turns a `metadata` Tensor into an index map function."""
+
+  def _index_map_func(index: Union[int, tensor.Tensor]) -> tuple[int, int]:
+    element_index = 0
+    while (element_index < metadata.shape[0] and
+           index >= array_ops.gather_nd(metadata, [element_index, 1])):
+      element_index += 1
+    offset = (
+        index - array_ops.gather_nd(metadata, [element_index, 0])
+        if element_index < metadata.shape[0]
+        else constant_op.constant(0, dtype=dtypes.int64))
+    return (element_index, offset)
+
+  return _index_map_func
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,230 @@
+# Copyright 2022 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.index_shuffle()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import shuffle_ops
+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.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class IndexShuffleTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _build_dataset(self,
+                     seed=None,
+                     reshuffle_each_iteration=None,
+                     num_elements=10):
+    file_infos = []
+    for _ in range(5):
+      file_infos.append({"path": "unused", "num_elements": num_elements})
+
+    def reader_factory(files):
+      return dataset_ops.Dataset.range(
+          num_elements * array_ops.shape(files, out_type=dtypes.int64)[0])
+
+    return shuffle_ops.index_shuffle(
+        file_infos,
+        reader_factory,
+        seed=seed,
+        reshuffle_each_iteration=reshuffle_each_iteration)
+
+  def testProcessFileInfos(self):
+    file_infos = []
+    file_infos.append({
+        "path": "take_50",
+        "num_elements": 100,
+        "skip": 25,
+        "take": 50
+    })
+    file_infos.append({
+        "path": "skip_all",
+        "num_elements": 100,
+        "skip": -1,
+    })
+    file_infos.append({
+        "path": "take_all",
+        "num_elements": 100,
+        "take": -1
+    })
+    file_infos.append({
+        "path": "take_10",
+        "num_elements": 100,
+        "skip": 90,
+        "take": 20
+    })
+    result = shuffle_ops._process_file_infos(file_infos)
+    self.assertEqual(result["files"],
+                     ["take_50", "skip_all", "take_all", "take_10"])
+    self.assertEqual(result["num_elements"], 160)
+    inputs = [0, 49, 50, 51, 149, 150, 151, 159]
+    expected = [25, 74, 200, 201, 299, 390, 391, 399]
+    for i, expected in enumerate(expected):
+      self.assertEqual(
+          self.evaluate(
+              shuffle_ops._adjust_index([inputs[i]], result["thresholds"],
+                                        result["offsets"])), expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnseeded(self):
+    # Assert that the shuffled dataset has the same elements as the
+    # "ground truth".
+    unshuffled_elements = np.arange(50)
+    shuffled_elements_1 = self.getDatasetOutput(
+        self._build_dataset(), requires_initialization=True)
+    shuffled_elements_2 = self.getDatasetOutput(
+        self._build_dataset(), requires_initialization=True)
+    self.assertAllEqual(
+        sorted(unshuffled_elements), sorted(shuffled_elements_1))
+    self.assertAllEqual(
+        sorted(unshuffled_elements), sorted(shuffled_elements_2))
+    self.assertNotEqual(shuffled_elements_1, shuffled_elements_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSameSeed(self):
+    # Assert that shuffling twice with the same seeds gives the same sequence.
+    shuffled_elements_1 = self.getDatasetOutput(
+        self._build_dataset(seed=42), requires_initialization=True)
+    shuffled_elements_2 = self.getDatasetOutput(
+        self._build_dataset(seed=42), requires_initialization=True)
+    self.assertEqual(shuffled_elements_1, shuffled_elements_2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLargeDataSet(self):
+    self._build_dataset(seed=42, num_elements=128*1024*1024)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentSeed(self):
+    # Assert that shuffling twice with a different seed gives a different
+    # permutation of the same elements.
+    shuffled_elements_1 = self.getDatasetOutput(
+        self._build_dataset(seed=42), requires_initialization=True)
+    shuffled_elements_2 = self.getDatasetOutput(
+        self._build_dataset(seed=24), requires_initialization=True)
+    self.assertNotEqual(shuffled_elements_1, shuffled_elements_2)
+    self.assertAllEqual(
+        sorted(shuffled_elements_1), sorted(shuffled_elements_2))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.v2_eager_only_combinations(),
+          combinations.combine(reshuffle_each_iteration=[True, False])))
+  def testReshuffleEachIteration(self, reshuffle_each_iteration):
+    # Assert that `reshuffle_each_iteration` controls whether data is shuffled
+    # differently across different epochs.
+    dataset = self._build_dataset(
+        seed=42, reshuffle_each_iteration=reshuffle_each_iteration)
+    shuffled_elements_1 = self.getDatasetOutput(dataset)
+    shuffled_elements_2 = self.getDatasetOutput(dataset)
+    if reshuffle_each_iteration:
+      self.assertNotEqual(shuffled_elements_1, shuffled_elements_2)
+      self.assertAllEqual(
+          sorted(shuffled_elements_1), sorted(shuffled_elements_2))
+    else:
+      self.assertAllEqual(shuffled_elements_1, shuffled_elements_2)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(skip=[8, 2, 0, -1], take=[8, 2, 0, -1])))
+  def testSkipAndTake(self, skip, take):
+    num_elements = 10
+    file_infos = []
+    file_infos.append({
+        "path": "unused",
+        "num_elements": num_elements,
+        "skip": skip if skip >= 0 else num_elements,
+        "take": take if take >= 0 else num_elements,
+    })
+    start = skip if skip >= 0 else num_elements
+    stop = min(num_elements, skip + take if take >= 0 else num_elements)
+    expected = np.arange(start, stop)
+
+    def reader_factory(_):
+      return dataset_ops.Dataset.range(10)
+
+    dataset = shuffle_ops.index_shuffle(file_infos, reader_factory)
+    actual = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertAllEqual(sorted(expected), sorted(actual))
+
+
+class IndexShuffleCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                 parameterized.TestCase):
+
+  def _build_dataset(
+      self,
+      num_elements_per_file,
+      num_files,
+      num_epochs,
+      seed=None,
+      reshuffle_each_iteration=None,
+      symbolic_checkpoint=None,
+  ):
+    file_infos = []
+    for _ in range(num_files):
+      file_infos.append({
+          "path": "unused",
+          "num_elements": num_elements_per_file,
+      })
+
+    def reader_factory(files):
+      return dataset_ops.Dataset.range(
+          num_elements_per_file *
+          array_ops.shape(files, out_type=dtypes.int64)[0])
+
+    dataset = shuffle_ops.index_shuffle(
+        file_infos,
+        reader_factory,
+        seed=seed,
+        reshuffle_each_iteration=reshuffle_each_iteration)
+    dataset = dataset.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=[False, True],
+              reshuffle_each_iteration=[False, True])))
+  def test(self, verify_fn, symbolic_checkpoint, reshuffle_each_iteration):
+    seed = 42
+    num_elements_per_file = 8
+    num_files = 3
+    num_epochs = 2
+    num_outputs = num_elements_per_file * num_files * num_epochs
+    # pylint: disable=g-long-lambda
+    verify_fn(
+        self, lambda: self._build_dataset(
+            num_elements_per_file=num_elements_per_file,
+            num_files=num_files,
+            num_epochs=num_epochs,
+            seed=seed,
+            reshuffle_each_iteration=reshuffle_each_iteration,
+            symbolic_checkpoint=symbolic_checkpoint), num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,191 @@
+# 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
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import io
+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(io.save(dataset, self._test_dir, compression=compression))
+    dataset2 = io.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(io.save(dataset, self._test_dir))
+    dataset2 = io.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(
+        io.save(dataset, self._test_dir, shard_func=lambda x: x // 21)
+    )
+    dataset2 = io.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(io.save(dataset, self._test_dir, shard_func=lambda x: x % 7))
+    dataset2 = io.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():
+      io.save(dataset, self._test_dir, compression=compression)
+
+    save_fn()
+    dataset = io.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(io.save(dataset, self._test_dir))
+    dataset_loaded = io.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(io.save(dataset1, self._test_dir))
+    dataset = io.load(self._test_dir)
+    dataset = dataset.shuffle(buffer_size=16)
+    dataset = dataset.batch(16)
+    dataset = dataset.repeat()
+    dataset = dataset.prefetch(1)
+    next_element = self.getNext(dataset)
+    for _ in range(30):
+      self.evaluate(next_element())
+
+
+class LoadCheckpointTest(IOTest, checkpoint_test_base.CheckpointTestBase):
+
+  def _build_ds(self):
+    return io.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(io.save(dataset, 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}
+    io.save(dataset, 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,
+    }
+    io.save(dataset, 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):
+      io.save(dataset, self._save_dir, checkpoint_args=checkpoint_args)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,103 @@
+# 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 tensorflow.python.data.experimental.ops.iterator_model_ops."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import iterator_model_ops
+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_impl
+from tensorflow.python.platform import test
+
+
+class IteratorModelOpsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def test_get_model_proto_not_empty(self):
+    options = dataset_ops.options_lib.Options()
+    options.autotune.enabled = True
+
+    dataset = dataset_ops.Dataset.range(100000)
+    dataset = dataset.map(lambda x: ("tf.data", "is the best ML data loader."))
+    dataset = dataset.batch(5, drop_remainder=True, num_parallel_calls=1)
+    dataset = dataset.with_options(options)
+
+    iterator = iter(dataset)
+    model_proto = iterator_model_ops.get_model_proto(iterator)
+    dataset_names = set(
+        model_proto.nodes[key].name for key in model_proto.nodes
+    )
+
+    self.assertNotEmpty(
+        dataset_names,
+        "The model proto from the iterator should contain at least 1"
+        " dataset op.",
+    )
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def test_get_model_proto_error_when_autotune_not_enabled(self):
+    options = dataset_ops.options_lib.Options()
+    options.autotune.enabled = False
+
+    dataset = dataset_ops.Dataset.range(100000)
+    dataset = dataset.map(lambda x: ("tf.data", "is the best ML data loader."))
+    dataset = dataset.batch(5, drop_remainder=True, num_parallel_calls=1)
+    dataset = dataset.with_options(options)
+
+    iterator = iter(dataset)
+
+    with self.assertRaisesRegex(
+        errors_impl.NotFoundError,
+        "Did you disable autotune",
+    ):
+      _ = iterator_model_ops.get_model_proto(iterator)
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def test_get_model_proto_from_numpy_iterator(self):
+    options = dataset_ops.options_lib.Options()
+    options.autotune.enabled = True
+
+    dataset = dataset_ops.Dataset.range(100000)
+    dataset = dataset.map(lambda x: ("tf.data", "is the best ML data loader."))
+    dataset = dataset.batch(5, drop_remainder=True, num_parallel_calls=1)
+    dataset = dataset.with_options(options)
+
+    iterator = dataset.as_numpy_iterator()
+
+    model_proto = iterator_model_ops.get_model_proto(iterator)
+    dataset_names = set(
+        model_proto.nodes[key].name for key in model_proto.nodes
+    )
+
+    self.assertNotEmpty(
+        dataset_names,
+        "The model proto from the iterator should contain at least 1"
+        " dataset op.",
+    )
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def test_get_model_proto_unsupported_in_graph_mode(self):
+    dataset = dataset_ops.Dataset.range(100000)
+
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.session():
+      with self.assertRaises(ValueError):
+        _ = iterator_model_ops.get_model_proto(iterator)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,152 @@
+# 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 tensorflow.python.data.experimental.ops.lookup_ops."""
+
+import os
+
+from tensorflow.python import tf2
+from tensorflow.python.data.experimental.ops import lookup_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers as reader_ops
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import lookup_ops as core_lookup_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class DatasetInitializerTest(test.TestCase):
+
+  def getHashTable(self):
+    if tf2.enabled():
+      return core_lookup_ops.StaticHashTable
+    else:
+      return core_lookup_ops.StaticHashTableV1
+
+  def initialize_table(self, table):
+    if not tf2.enabled():
+      self.evaluate(table.initializer)
+
+  def _createVocabFile(self, basename, values=("brain", "salad", "surgery")):
+    vocabulary_file = os.path.join(self.get_temp_dir(), basename)
+    with open(vocabulary_file, "w") as f:
+      f.write("\n".join(values) + "\n")
+    return vocabulary_file
+
+  def test_basic(self):
+    keys = dataset_ops.Dataset.range(100)
+    values = dataset_ops.Dataset.range(100).map(
+        lambda x: string_ops.as_string(x * 2))
+    ds = dataset_ops.Dataset.zip((keys, values))
+    init = lookup_ops.DatasetInitializer(ds)
+    table = self.getHashTable()(init, default_value="")
+    self.initialize_table(table)
+
+    output = table.lookup(constant_op.constant([0, 2, 5], dtypes.int64))
+    result = self.evaluate(output)
+    self.assertAllEqual(["0", "4", "10"], result)
+
+  def test_basic_bad_shape(self):
+    keys = dataset_ops.Dataset.range(100)
+    values = dataset_ops.Dataset.range(100).map(
+        lambda x: string_ops.as_string(x * 2))
+    values = values.batch(4)
+    ds = dataset_ops.Dataset.zip((keys, values))
+    with self.assertRaises(ValueError):
+      lookup_ops.DatasetInitializer(ds)
+
+  def test_from_file(self):
+    vocabulary_file = self._createVocabFile("test.txt", ("one", "two", "three"))
+    ds = reader_ops.TextLineDataset(vocabulary_file)
+    ds = ds.enumerate(start=1)
+    init = lookup_ops.DatasetInitializer(ds)
+    table = self.getHashTable()(init, default_value="")
+    self.initialize_table(table)
+
+    output = table.lookup(constant_op.constant([2, 3, 4], dtypes.int64))
+    result = self.evaluate(output)
+    self.assertAllEqual(["two", "three", ""], result)
+
+  def test_from_multiple_files(self):
+    vocabulary_file1 = self._createVocabFile("test1.txt",
+                                             ("one", "two", "three"))
+    vocabulary_file2 = self._createVocabFile("test2.txt",
+                                             ("four", "five", "six"))
+    ds = reader_ops.TextLineDataset([vocabulary_file1, vocabulary_file2])
+    ds = ds.enumerate(start=1)
+    init = lookup_ops.DatasetInitializer(ds)
+    table = self.getHashTable()(init, default_value="")
+    self.initialize_table(table)
+
+    output = table.lookup(constant_op.constant([2, 3, 4], dtypes.int64))
+    result = self.evaluate(output)
+    self.assertAllEqual(["two", "three", "four"], result)
+
+  def test_map_variable(self):
+    ds = dataset_ops.Dataset.range(100)
+    captured_var = variables.Variable(0)
+
+    def func(_):
+      return captured_var.assign_add(1)
+
+    ds = ds.map(func)
+    ds = ds.enumerate(start=1)
+    init = lookup_ops.DatasetInitializer(ds)
+    table = self.getHashTable()(init, default_value=-1)
+    self.evaluate(captured_var.initializer)
+    self.initialize_table(table)
+
+    output = table.lookup(constant_op.constant([1, 2, 101], dtypes.int64))
+    result = self.evaluate(output)
+    self.assertAllEqual([1, 2, -1], result)
+
+  def test_compatibility(self):
+    with ops.Graph().as_default():
+      keys = dataset_ops.Dataset.range(100)
+      values = dataset_ops.Dataset.range(100).map(string_ops.as_string)
+      ds = dataset_ops.Dataset.zip((keys, values))
+      init = lookup_ops.DatasetInitializer(ds)
+      table = self.getHashTable()(init, default_value="")
+      output = table.lookup(constant_op.constant([0, 2, 5], dtypes.int64))
+      self.evaluate(core_lookup_ops.tables_initializer())
+      result = self.evaluate(output)
+    self.assertAllEqual(["0", "2", "5"], result)
+
+  def test_table_from_dataset(self):
+    keys = dataset_ops.Dataset.from_tensor_slices([2, 3, 4])
+    values = dataset_ops.Dataset.from_tensor_slices(["two", "three", "four"])
+    ds = dataset_ops.Dataset.zip((keys, values))
+    table = lookup_ops.table_from_dataset(
+        ds, default_value="n/a", key_dtype=dtypes.int64)
+    output = table.lookup(constant_op.constant([2, 3, 4], dtypes.int32))
+    self.evaluate(core_lookup_ops.tables_initializer())
+    result = self.evaluate(output)
+    self.assertAllEqual(["two", "three", "four"], result)
+
+  def test_index_table_from_dataset(self):
+    ds = dataset_ops.Dataset.range(100).map(
+        lambda x: string_ops.as_string(x * 2))
+    table = lookup_ops.index_table_from_dataset(ds, key_dtype=dtypes.int64)
+    output = table.lookup(
+        constant_op.constant(["0", "2", "4"], dtype=dtypes.string))
+    self.evaluate(core_lookup_ops.tables_initializer())
+    result = self.evaluate(output)
+    self.assertAllEqual([0, 1, 2], result)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,250 @@
+# 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.make_batched_features_dataset()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import readers
+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 readers as core_readers
+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 tensor as tensor_lib
+from tensorflow.python.ops import io_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.platform import test
+
+
+class MakeBatchedFeaturesDatasetTest(tf_record_test_base.FeaturesTestBase,
+                                     parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(batch_size=[1, 2], num_epochs=[1, 10]),
+          test_base.default_test_combinations()))
+  def testRead(self, batch_size, num_epochs):
+    # Basic test: read from file 0.
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[0],
+            label_key="label",
+            num_epochs=num_epochs,
+            batch_size=batch_size))
+    self._verify_records(
+        batch_size, 0, num_epochs=num_epochs, label_key_provided=True)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch(label_key_provided=True)
+
+      # Basic test: read from file 1.
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[1],
+            label_key="label",
+            num_epochs=num_epochs,
+            batch_size=batch_size))
+    self._verify_records(
+        batch_size, 1, num_epochs=num_epochs, label_key_provided=True)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch(label_key_provided=True)
+
+    # Basic test: read from both files.
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames,
+            label_key="label",
+            num_epochs=num_epochs,
+            batch_size=batch_size))
+    self._verify_records(
+        batch_size, num_epochs=num_epochs, label_key_provided=True)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch(label_key_provided=True)
+    # Basic test: read from both files.
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames,
+            num_epochs=num_epochs,
+            batch_size=batch_size))
+    self._verify_records(batch_size, num_epochs=num_epochs)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadWithEquivalentDataset(self):
+    features = {
+        "file": parsing_ops.FixedLenFeature([], dtypes.int64),
+        "record": parsing_ops.FixedLenFeature([], dtypes.int64),
+    }
+    dataset = (
+        core_readers.TFRecordDataset(self._filenames)
+        .map(lambda x: parsing_ops.parse_single_example(x, features))
+        .repeat(10).batch(2))
+    next_element = self.getNext(dataset)
+    for file_batch, _, _, _, record_batch, _ in self._next_expected_batch(
+        range(self._num_files), 2, 10):
+      actual_batch = self.evaluate(next_element())
+      self.assertAllEqual(file_batch, actual_batch["file"])
+      self.assertAllEqual(record_batch, actual_batch["record"])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(batch_size=[1, 2], num_epochs=[5]),
+          test_base.default_test_combinations()))
+  def testReadWithFusedShuffleRepeatDatasetSameSeed(self, batch_size,
+                                                    num_epochs):
+    total_records = num_epochs * self._num_records
+    # Test that shuffling with same seed produces the same result.
+    outputs1 = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[0],
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            shuffle=True,
+            shuffle_seed=5))
+    outputs2 = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[0],
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            shuffle=True,
+            shuffle_seed=5))
+    for _ in range(total_records // batch_size):
+      batch1 = self._run_actual_batch(outputs1)
+      batch2 = self._run_actual_batch(outputs2)
+      for i in range(len(batch1)):
+        self.assertAllEqual(batch1[i], batch2[i])
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(batch_size=[1, 2], num_epochs=[5]),
+          test_base.default_test_combinations()))
+  def testReadWithFusedShuffleRepeatDatasetDifferentSeed(
+      self, batch_size, num_epochs):
+    total_records = num_epochs * self._num_records
+    # Test that shuffling with different seeds produces a different order.
+    outputs1 = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[0],
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            shuffle=True,
+            shuffle_seed=5))
+    outputs2 = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames[0],
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            shuffle=True,
+            shuffle_seed=15))
+    all_equal = True
+    for _ in range(total_records // batch_size):
+      batch1 = self._run_actual_batch(outputs1)
+      batch2 = self._run_actual_batch(outputs2)
+      for i in range(len(batch1)):
+        all_equal = all_equal and np.array_equal(batch1[i], batch2[i])
+    self.assertFalse(all_equal)
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(
+              batch_size=[1, 2],
+              num_epochs=[5],
+              reader_num_threads=[2, 4],
+              parser_num_threads=[2, 4]),
+          test_base.default_test_combinations()))
+  def testParallelReadersAndParsers(self, batch_size, num_epochs,
+                                    reader_num_threads, parser_num_threads):
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames,
+            label_key="label",
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            reader_num_threads=reader_num_threads,
+            parser_num_threads=parser_num_threads))
+    self._verify_records(
+        batch_size,
+        num_epochs=num_epochs,
+        label_key_provided=True,
+        interleave_cycle_length=reader_num_threads)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch(label_key_provided=True)
+
+    self.outputs = self.getNext(
+        self.make_batch_feature(
+            filenames=self._filenames,
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            reader_num_threads=reader_num_threads,
+            parser_num_threads=parser_num_threads))
+    self._verify_records(
+        batch_size,
+        num_epochs=num_epochs,
+        interleave_cycle_length=reader_num_threads)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._next_actual_batch()
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(batch_size=[1, 2], num_epochs=[1, 10]),
+          test_base.default_test_combinations()))
+  def testDropFinalBatch(self, batch_size, num_epochs):
+    # Basic test: read from file 0.
+    outputs = self.make_batch_feature(
+        filenames=self._filenames[0],
+        label_key="label",
+        num_epochs=num_epochs,
+        batch_size=batch_size,
+        drop_final_batch=True)
+    for tensor in nest.flatten(outputs):
+      if isinstance(tensor, tensor_lib.Tensor):  # Guard against SparseTensor.
+        self.assertEqual(tensor.shape[0], batch_size)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIndefiniteRepeatShapeInference(self):
+    dataset = self.make_batch_feature(
+        filenames=self._filenames[0],
+        label_key="label",
+        num_epochs=None,
+        batch_size=32)
+    for shape, clazz in zip(
+        nest.flatten(dataset_ops.get_legacy_output_shapes(dataset)),
+        nest.flatten(dataset_ops.get_legacy_output_classes(dataset))):
+      if issubclass(clazz, tensor_lib.Tensor):
+        self.assertEqual(32, shape[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOldStyleReader(self):
+    with self.assertRaisesRegex(
+        TypeError, r"The `reader` argument must return a `Dataset` object. "
+        r"`tf.ReaderBase` subclasses are not supported."):
+      _ = readers.make_batched_features_dataset(
+          file_pattern=self._filenames[0], batch_size=32,
+          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),
+          },
+          reader=io_ops.TFRecordReader)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,793 @@
+# 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.make_csv_dataset()`."""
+import gzip
+import os
+import zlib
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import readers
+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 constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class MakeCsvDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _make_csv_dataset(self, filenames, batch_size, num_epochs=1, **kwargs):
+    return readers.make_csv_dataset(
+        filenames, batch_size=batch_size, num_epochs=num_epochs, **kwargs)
+
+  def _setup_files(self,
+                   inputs,
+                   linebreak="\n",
+                   compression_type=None,
+                   encoding="utf-8"):
+    filenames = []
+    for i, ip in enumerate(inputs):
+      fn = os.path.join(self.get_temp_dir(), "temp_%d.csv" % i)
+      contents = linebreak.join(ip).encode(encoding)
+      if compression_type is None:
+        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)
+      filenames.append(fn)
+    return filenames
+
+  def _next_expected_batch(self, expected_output, expected_keys, batch_size,
+                           num_epochs):
+    features = {k: [] for k in expected_keys}
+    for _ in range(num_epochs):
+      for values in expected_output:
+        for n, key in enumerate(expected_keys):
+          features[key].append(values[n])
+        if len(features[expected_keys[0]]) == batch_size:
+          yield features
+          features = {k: [] for k in expected_keys}
+    if features[expected_keys[0]]:  # Leftover from the last batch
+      yield features
+
+  def _verify_output(
+      self,
+      dataset,
+      batch_size,
+      num_epochs,
+      label_name,
+      expected_output,
+      expected_keys,
+  ):
+    get_next = self.getNext(dataset)
+
+    for expected_features in self._next_expected_batch(
+        expected_output,
+        expected_keys,
+        batch_size,
+        num_epochs,
+    ):
+      actual_features = self.evaluate(get_next())
+
+      if label_name is not None:
+        expected_labels = expected_features.pop(label_name)
+        self.assertAllEqual(expected_labels, actual_features[1])
+        actual_features = actual_features[0]
+
+      for k in expected_features.keys():
+        # Compare features
+        self.assertAllEqual(expected_features[k], actual_features[k])
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  def _test_dataset(self,
+                    inputs,
+                    expected_output,
+                    expected_keys,
+                    batch_size=1,
+                    num_epochs=1,
+                    label_name=None,
+                    encoding="utf-8",
+                    **kwargs):
+    """Checks that elements produced by CsvDataset match expected output."""
+    # Convert str type because py3 tf strings are bytestrings
+    filenames = self._setup_files(
+        inputs,
+        compression_type=kwargs.get("compression_type", None),
+        encoding=encoding)
+    dataset = self._make_csv_dataset(
+        filenames,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        label_name=label_name,
+        encoding=encoding,
+        **kwargs)
+    self._verify_output(dataset, batch_size, num_epochs, label_name,
+                        expected_output, expected_keys)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    """Tests making a CSV dataset with keys and defaults provided."""
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEncoding(self):
+    """Tests making a CSV dataset with an encoding except for utf-8."""
+    record_defaults = [
+        constant_op.constant([], dtypes.string),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(2)]
+    inputs = [[",".join(x for x in column_names), "さる,猿", "とり,鳥"],
+              [",".join(x for x in column_names), "いぬ,犬", "ねこ,猫"]]
+    expected_output = [["さる".encode("shift-jis"), "猿".encode("shift-jis")],
+                       ["とり".encode("shift-jis"), "鳥".encode("shift-jis")],
+                       ["いぬ".encode("shift-jis"), "犬".encode("shift-jis")],
+                       ["ねこ".encode("shift-jis"), "猫".encode("shift-jis")]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+        encoding="shift-jis",
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithBatchSizeAndEpochs(self):
+    """Tests making a CSV dataset with keys and defaults provided."""
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=3,
+        num_epochs=10,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCompressionType(self):
+    """Tests `compression_type` argument."""
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    for compression_type in ("GZIP", "ZLIB"):
+      self._test_dataset(
+          inputs,
+          expected_output=expected_output,
+          expected_keys=column_names,
+          column_names=column_names,
+          label_name=label,
+          batch_size=1,
+          num_epochs=1,
+          shuffle=False,
+          header=True,
+          column_defaults=record_defaults,
+          compression_type=compression_type,
+      )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithCompressionTypeAndNoColumnNames(self):
+    """Tests `compression_type` argument."""
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"],
+              [
+                  ",".join(x for x in column_names), "10,11,12,13,14",
+                  "15,16,17,18,19"
+              ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+        compression_type="GZIP",
+    )
+
+    with self.assertRaisesRegex(ValueError,
+                                "`compression_type` ZLIB is not supported"):
+      self._test_dataset(
+          inputs,
+          expected_output=expected_output,
+          expected_keys=column_names,
+          label_name=label,
+          batch_size=1,
+          num_epochs=1,
+          shuffle=False,
+          header=True,
+          column_defaults=record_defaults,
+          compression_type="ZLIB",
+      )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithBadInputs(self):
+    """Tests that exception is raised when input is malformed.
+    """
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    filenames = self._setup_files(inputs)
+
+    # Duplicate column names
+    with self.assertRaises(ValueError):
+      self._make_csv_dataset(
+          filenames,
+          batch_size=1,
+          column_defaults=record_defaults,
+          label_name="col0",
+          column_names=column_names * 2)
+
+    # Label key not one of column names
+    with self.assertRaises(ValueError):
+      self._make_csv_dataset(
+          filenames,
+          batch_size=1,
+          column_defaults=record_defaults,
+          label_name="not_a_real_label",
+          column_names=column_names)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNoLabel(self):
+    """Tests making a CSV dataset with no label provided."""
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNoHeader(self):
+    """Tests that datasets can be created from CSV files with no header line.
+    """
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [["0,1,2,3,4", "5,6,7,8,9"], ["10,11,12,13,14", "15,16,17,18,19"]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=False,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithTypes(self):
+    """Tests that defaults can be a dtype instead of a Tensor for required vals.
+    """
+    record_defaults = [
+        dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64,
+        dtypes.string
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x[0] for x in column_names), "0,1,2,3,4", "5,6,7,8,9"],
+              [
+                  ",".join(x[0] for x in column_names), "10,11,12,13,14",
+                  "15,16,17,18,19"
+              ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNoColNames(self):
+    """Tests that datasets can be created when column names are not specified.
+
+    In that case, we should infer the column names from the header lines.
+    """
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    expected_output = [[0, 1, 2, 3, b"4"], [5, 6, 7, 8, b"9"],
+                       [10, 11, 12, 13, b"14"], [15, 16, 17, 18, b"19"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        column_defaults=record_defaults,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithTypeInferenceMismatch(self):
+    # Test that error is thrown when num fields doesn't match columns
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    filenames = self._setup_files(inputs)
+    with self.assertRaises(ValueError):
+      self._make_csv_dataset(
+          filenames,
+          column_names=column_names + ["extra_name"],
+          column_defaults=None,
+          batch_size=2,
+          num_epochs=10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithTypeInference(self):
+    """Tests that datasets can be created when no defaults are specified.
+
+    In that case, we should infer the types from the first N records.
+    """
+    column_names = ["col%d" % i for i in range(5)]
+    str_int32_max = str(2**33)
+    inputs = [[
+        ",".join(x for x in column_names),
+        "0,%s,2.0,3e50,rabbit" % str_int32_max
+    ]]
+    expected_output = [[0, 2**33, 2.0, 3e50, b"rabbit"]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithTypeInferenceFallthrough(self):
+    """Tests that datasets can be created when no defaults are specified.
+
+    Tests on a deliberately tricky file.
+    """
+    column_names = ["col%d" % i for i in range(5)]
+    str_int32_max = str(2**33)
+    inputs = [[
+        ",".join(x for x in column_names),
+        ",,,,",
+        "0,0,0.0,0.0,0.0",
+        "0,%s,2.0,3e50,rabbit" % str_int32_max,
+        ",,,,",
+    ]]
+    expected_output = [[0, 0, 0, 0, b""], [0, 0, 0, 0, b"0.0"],
+                       [0, 2**33, 2.0, 3e50, b"rabbit"], [0, 0, 0, 0, b""]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithNAValuesAndFieldDelim(self):
+    """Tests that datasets can be created from different delim and na_value."""
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [["0 1 2 3 4", "5 6 7 8 9"], ["10 11 12 13 14", "15 16 17 ? 19"]]
+    expected_output = [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14],
+                       [15, 16, 17, 0, 19]]
+    label = "col0"
+
+    self._test_dataset(
+        inputs,
+        expected_output=expected_output,
+        expected_keys=column_names,
+        column_names=column_names,
+        label_name=label,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=False,
+        na_value="?",
+        field_delim=" ",
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithSelectCols(self):
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+    column_names = ["col%d" % i for i in range(5)]
+    str_int32_max = str(2**33)
+    inputs = [[
+        ",".join(x for x in column_names),
+        "0,%s,2.0,3e50,rabbit" % str_int32_max
+    ]]
+    expected_output = [[0, 2**33, 2.0, 3e50, b"rabbit"]]
+
+    select_cols = [1, 3, 4]
+    self._test_dataset(
+        inputs,
+        expected_output=[[x[i] for i in select_cols] for x in expected_output],
+        expected_keys=[column_names[i] for i in select_cols],
+        column_names=column_names,
+        column_defaults=[record_defaults[i] for i in select_cols],
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        select_columns=select_cols,
+    )
+
+    # Can still do inference without provided defaults
+    self._test_dataset(
+        inputs,
+        expected_output=[[x[i] for i in select_cols] for x in expected_output],
+        expected_keys=[column_names[i] for i in select_cols],
+        column_names=column_names,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        select_columns=select_cols,
+    )
+
+    # Can still do column name inference
+    self._test_dataset(
+        inputs,
+        expected_output=[[x[i] for i in select_cols] for x in expected_output],
+        expected_keys=[column_names[i] for i in select_cols],
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        select_columns=select_cols,
+    )
+
+    # Can specify column names instead of indices
+    self._test_dataset(
+        inputs,
+        expected_output=[[x[i] for i in select_cols] for x in expected_output],
+        expected_keys=[column_names[i] for i in select_cols],
+        column_names=column_names,
+        batch_size=1,
+        num_epochs=1,
+        shuffle=False,
+        header=True,
+        select_columns=[column_names[i] for i in select_cols],
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithSelectColsError(self):
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+    column_names = ["col%d" % i for i in range(5)]
+    str_int32_max = str(2**33)
+    inputs = [[
+        ",".join(x for x in column_names),
+        "0,%s,2.0,3e50,rabbit" % str_int32_max
+    ]]
+
+    select_cols = [1, 3, 4]
+    filenames = self._setup_files(inputs)
+
+    with self.assertRaises(ValueError):
+      # Mismatch in number of defaults and number of columns selected,
+      # should raise an error
+      self._make_csv_dataset(
+          filenames,
+          batch_size=1,
+          column_defaults=record_defaults,
+          column_names=column_names,
+          select_columns=select_cols)
+
+    with self.assertRaises(ValueError):
+      # Invalid column name should raise an error
+      self._make_csv_dataset(
+          filenames,
+          batch_size=1,
+          column_defaults=[[0]],
+          column_names=column_names,
+          label_name=None,
+          select_columns=["invalid_col_name"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithShuffle(self):
+    record_defaults = [
+        constant_op.constant([], dtypes.int32),
+        constant_op.constant([], dtypes.int64),
+        constant_op.constant([], dtypes.float32),
+        constant_op.constant([], dtypes.float64),
+        constant_op.constant([], dtypes.string)
+    ]
+
+    def str_series(st):
+      return ",".join(str(i) for i in range(st, st + 5))
+
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [
+        [",".join(x for x in column_names)
+        ] + [str_series(5 * i) for i in range(15)],
+        [",".join(x for x in column_names)] +
+        [str_series(5 * i) for i in range(15, 20)],
+    ]
+
+    filenames = self._setup_files(inputs)
+
+    total_records = 20
+    for batch_size in [1, 2]:
+      # Test that shuffling with the same seed produces the same result
+      dataset1 = self._make_csv_dataset(
+          filenames,
+          column_defaults=record_defaults,
+          column_names=column_names,
+          batch_size=batch_size,
+          header=True,
+          shuffle=True,
+          shuffle_seed=5,
+          num_epochs=2,
+      )
+      dataset2 = self._make_csv_dataset(
+          filenames,
+          column_defaults=record_defaults,
+          column_names=column_names,
+          batch_size=batch_size,
+          header=True,
+          shuffle=True,
+          shuffle_seed=5,
+          num_epochs=2,
+      )
+      next1 = self.getNext(dataset1)
+      next2 = self.getNext(dataset2)
+      for _ in range(total_records // batch_size):
+        batch1 = nest.flatten(self.evaluate(next1()))
+        batch2 = nest.flatten(self.evaluate(next2()))
+        for i in range(len(batch1)):
+          self.assertAllEqual(batch1[i], batch2[i])
+
+      # Test that shuffling with a different seed produces different results
+      dataset1 = self._make_csv_dataset(
+          filenames,
+          column_defaults=record_defaults,
+          column_names=column_names,
+          batch_size=batch_size,
+          header=True,
+          shuffle=True,
+          shuffle_seed=5,
+          num_epochs=2,
+      )
+      dataset2 = self._make_csv_dataset(
+          filenames,
+          column_defaults=record_defaults,
+          column_names=column_names,
+          batch_size=batch_size,
+          header=True,
+          shuffle=True,
+          shuffle_seed=6,
+          num_epochs=2,
+      )
+      next1 = self.getNext(dataset1)
+      next2 = self.getNext(dataset2)
+      all_equal = False
+      for _ in range(total_records // batch_size):
+        batch1 = nest.flatten(self.evaluate(next1()))
+        batch2 = nest.flatten(self.evaluate(next2()))
+        for i in range(len(batch1)):
+          all_equal = all_equal and np.array_equal(batch1[i], batch2[i])
+      self.assertFalse(all_equal)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIndefiniteRepeatShapeInference(self):
+    column_names = ["col%d" % i for i in range(5)]
+    inputs = [[",".join(x for x in column_names), "0,1,2,3,4", "5,6,7,8,9"], [
+        ",".join(x for x in column_names), "10,11,12,13,14", "15,16,17,18,19"
+    ]]
+    filenames = self._setup_files(inputs)
+    dataset = self._make_csv_dataset(filenames, batch_size=32, num_epochs=None)
+    for shape in nest.flatten(dataset_ops.get_legacy_output_shapes(dataset)):
+      self.assertEqual(32, shape[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFieldOrder(self):
+    data = [[
+        "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19",
+        "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19"
+    ]]
+    file_path = self._setup_files(data)
+
+    ds = readers.make_csv_dataset(
+        file_path, batch_size=1, shuffle=False, num_epochs=1)
+    nxt = self.getNext(ds)
+
+    result = list(self.evaluate(nxt()).values())
+
+    self.assertEqual(result, sorted(result))
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,97 @@
+# 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.
+# ==============================================================================
+"""Integration test for dataset serialization."""
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import iterator_ops as contrib_iterator_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+
+
+class MakeSaveableFromIteratorTest(test.TestCase, parameterized.TestCase):
+
+  def _build_input_pipeline(self, name, num_outputs):
+    with ops.name_scope(name):
+      ds = dataset_ops.Dataset.range(num_outputs).shuffle(
+          10, reshuffle_each_iteration=False).prefetch(10)
+      iterator = ds.make_initializable_iterator()
+      saveable = contrib_iterator_ops.make_saveable_from_iterator(iterator)
+      ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable)
+      return iterator.initializer, iterator.get_next()
+
+  def _build_graph(self, num_pipelines, num_outputs):
+    init_ops = []
+    get_next_ops = []
+    for i in range(num_pipelines):
+      name = "input_pipeline_%d" % i
+      init_op, get_next_op = self._build_input_pipeline(name, num_outputs)
+      init_ops.append(init_op)
+      get_next_ops.append(get_next_op)
+    saver = saver_lib.Saver()
+    return init_ops, get_next_ops, saver
+
+  def _ckpt_path(self):
+    return os.path.join(self.get_temp_dir(), "iterator")
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testConcurrentSaves(self):
+    num_pipelines = 10
+    num_outputs = 10
+    break_point = 10
+    all_outputs = [[] for _ in range(num_pipelines)]
+    with ops.Graph().as_default() as g:
+      init_ops, get_next_ops, saver = self._build_graph(num_pipelines,
+                                                        num_outputs)
+      with self.session(graph=g) as sess:
+        self.evaluate(init_ops)
+        for _ in range(break_point):
+          output = self.evaluate(get_next_ops)
+          for i in range(num_pipelines):
+            all_outputs[i].append(output[i])
+        saver.save(sess, self._ckpt_path())
+
+    with ops.Graph().as_default() as g:
+      init_ops, get_next_ops, saver = self._build_graph(num_pipelines,
+                                                        num_outputs)
+      with self.session(graph=g) as sess:
+        self.evaluate(init_ops)
+        saver.restore(sess, self._ckpt_path())
+        for _ in range(num_outputs - break_point):
+          output = self.evaluate(get_next_ops)
+          for i in range(num_pipelines):
+            all_outputs[i].append(output[i])
+
+    for output in all_outputs:
+      self.assertSequenceEqual(sorted(output), range(num_outputs))
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testUninitializedIterator(self):
+    num_pipelines = 1
+    num_outputs = 1
+    with ops.Graph().as_default() as g:
+      _, _, saver = self._build_graph(num_pipelines, num_outputs)
+      with self.session(graph=g) as sess:
+        with self.assertRaises(errors.FailedPreconditionError):
+          saver.save(sess, self._ckpt_path())
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,135 @@
+# 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.make_tf_record_dataset()`."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import readers
+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.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+
+
+class MakeTFRecordDatasetTest(tf_record_test_base.TFRecordTestBase,
+                              parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              batch_size=[1, 2],
+              num_epochs=[1, 3],
+              file_index=[None, 1],
+              num_parallel_reads=[1, 8],
+              drop_final_batch=[False, True],
+              parser_fn=[True, False])))
+  def testRead(self, batch_size, num_epochs, file_index, num_parallel_reads,
+               drop_final_batch, parser_fn):
+    if file_index is None:
+      file_pattern = self._filenames
+    else:
+      file_pattern = self._filenames[file_index]
+
+    if parser_fn:
+      fn = lambda x: string_ops.substr(x, 1, 999)
+    else:
+      fn = None
+
+    outputs = self.getNext(
+        readers.make_tf_record_dataset(
+            file_pattern=file_pattern,
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            parser_fn=fn,
+            num_parallel_reads=num_parallel_reads,
+            drop_final_batch=drop_final_batch,
+            shuffle=False))
+    self._verify_records(
+        outputs,
+        batch_size,
+        file_index,
+        num_epochs=num_epochs,
+        interleave_cycle_length=num_parallel_reads,
+        drop_final_batch=drop_final_batch,
+        use_parser_fn=parser_fn)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(outputs())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              batch_size=[1, 2],
+              num_epochs=[1, 3],
+              num_parallel_reads=[1, 2],
+              seed=[None, 21345])))
+  def testShuffle(self, batch_size, num_epochs, num_parallel_reads, seed):
+
+    def dataset_fn():
+      return readers.make_tf_record_dataset(
+          file_pattern=self._filenames,
+          num_epochs=num_epochs,
+          batch_size=batch_size,
+          num_parallel_reads=num_parallel_reads,
+          shuffle=True,
+          shuffle_seed=seed)
+
+    next_element = self.getNext(dataset_fn())
+    first_batches = []
+    try:
+      while True:
+        first_batches.append(self.evaluate(next_element()))
+    except errors.OutOfRangeError:
+      pass
+
+    next_element = self.getNext(dataset_fn())
+    second_batches = []
+    try:
+      while True:
+        second_batches.append(self.evaluate(next_element()))
+    except errors.OutOfRangeError:
+      pass
+
+    self.assertEqual(len(first_batches), len(second_batches))
+    if seed is not None:
+      # if you set a seed, should get the same results
+      for i in range(len(first_batches)):
+        self.assertAllEqual(first_batches[i], second_batches[i])
+
+    expected = []
+    for f in range(self._num_files):
+      for r in range(self._num_records):
+        expected.extend([self._record(f, r)] * num_epochs)
+
+    for batches in (first_batches, second_batches):
+      actual = []
+      for b in batches:
+        actual.extend(b)
+      self.assertAllEqual(sorted(expected), sorted(actual))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIndefiniteRepeatShapeInference(self):
+    dataset = readers.make_tf_record_dataset(
+        file_pattern=self._filenames, num_epochs=None, batch_size=32)
+    for shape in nest.flatten(dataset_ops.get_legacy_output_shapes(dataset)):
+      self.assertEqual(32, shape[0])
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,517 @@
+# 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.map_and_batch()`."""
+import math
+
+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 batching
+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 sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class MapAndBatchTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_parallel_calls=[None, 1, 2], num_parallel_batches=None) +
+          combinations.combine(
+              num_parallel_calls=None, num_parallel_batches=10)))
+  def testMapAndBatch(self, num_parallel_calls, num_parallel_batches):
+    """Test a dataset that maps a TF function across its input elements."""
+    # The pipeline is TensorSliceDataset ->
+    # RepeatDataset(count) -> MapAndBatchDataset(square_3, 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)
+
+    def dataset_fn(batch_size, count):
+      dataset = dataset_ops.Dataset.from_tensor_slices(components).repeat(
+          count).apply(
+              batching.map_and_batch(
+                  map_func=_map_fn,
+                  batch_size=batch_size,
+                  num_parallel_calls=num_parallel_calls,
+                  num_parallel_batches=num_parallel_batches))
+      return dataset
+
+    # Batch of a finite input, where the batch_size divides the
+    # total number of elements.
+    dataset = dataset_fn(14, 28)
+    get_next = self.getNext(dataset)
+    self.assertEqual(
+        [[None] + list(c.shape[1:]) for c in components],
+        [shape.as_list()
+         for shape in dataset_ops.get_legacy_output_shapes(dataset)])
+    num_batches = (28 * 7) // 14
+    for i in range(num_batches):
+      result = self.evaluate(get_next())
+      for component, result_component in zip(components, result):
+        for j in range(14):
+          self.assertAllEqual(component[(i * 14 + j) % 7]**2,
+                              result_component[j])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Batch of a finite input, where the batch_size does not
+    # divide the total number of elements.
+    get_next = self.getNext(dataset_fn(8, 14))
+
+    # We expect (num_batches - 1) full-sized batches.
+    num_batches = int(math.ceil((14 * 7) / 8))
+    for i in range(num_batches - 1):
+      result = self.evaluate(get_next())
+      for component, result_component in zip(components, result):
+        for j in range(8):
+          self.assertAllEqual(component[(i * 8 + j) % 7]**2,
+                              result_component[j])
+
+    result = self.evaluate(get_next())
+    for component, result_component in zip(components, result):
+      for j in range((14 * 7) % 8):
+        self.assertAllEqual(component[((num_batches - 1) * 8 + j) % 7]**2,
+                            result_component[j])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+    # Batch of an empty input should fail straight away.
+    self.assertDatasetProduces(dataset_fn(8, 0), expected_output=[])
+
+    # Empty batch should be an initialization time error.
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.assertDatasetProduces(dataset_fn(0, 14), expected_output=[])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testMapAndBatchPartialBatch(self, drop_remainder):
+    dataset = (
+        dataset_ops.Dataset.range(10).apply(
+            batching.map_and_batch(
+                lambda x: array_ops.reshape(x * x, [1]),
+                batch_size=4,
+                drop_remainder=drop_remainder)))
+
+    if drop_remainder:
+      self.assertEqual(
+          [4, 1], dataset_ops.get_legacy_output_shapes(dataset).as_list())
+    else:
+      self.assertEqual(
+          [None, 1], dataset_ops.get_legacy_output_shapes(dataset).as_list())
+    expected_output = [[[0], [1], [4], [9]], [[16], [25], [36], [49]]]
+    if not drop_remainder:
+      expected_output.append([[64], [81]])
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchYieldsPartialBatch(self):
+    dataset = (
+        dataset_ops.Dataset.range(10).apply(
+            batching.map_and_batch(lambda x: array_ops.reshape(x * x, [1]), 4)))
+
+    self.assertEqual(
+        [None, 1], dataset_ops.get_legacy_output_shapes(dataset).as_list())
+    expected_output = [[[0], [1], [4], [9]], [[16], [25], [36], [49]],
+                       [[64], [81]]]
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchParallelGetNext(self):
+    dataset = dataset_ops.Dataset.range(50000).apply(
+        batching.map_and_batch(lambda x: x, batch_size=100))
+
+    if context.executing_eagerly():
+      iterator = iter(dataset)
+      get_next = iterator._next_internal  # pylint: disable=protected-access
+    else:
+      iterator = dataset_ops.make_one_shot_iterator(dataset)
+      get_next = iterator.get_next
+
+    elements = []
+    for _ in range(100):
+      elements.append(get_next)
+
+    for i in range(5):
+      got = self.evaluate([element() for element in elements])
+      got.sort(key=lambda x: x[0])
+      expected = []
+      for j in range(100):
+        expected.append(range(i * 10000 + j * 100, i * 10000 + (j + 1) * 100))
+      self.assertAllEqual(got, expected)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate([element() for element in elements])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchParallelGetNextDropRemainder(self):
+    dataset = dataset_ops.Dataset.range(49999).apply(
+        batching.map_and_batch(
+            lambda x: x, batch_size=100, drop_remainder=True))
+
+    if context.executing_eagerly():
+      iterator = iter(dataset)
+      get_next = iterator._next_internal  # pylint: disable=protected-access
+    else:
+      iterator = dataset_ops.make_one_shot_iterator(dataset)
+      get_next = iterator.get_next
+
+    elements = []
+    for _ in range(100):
+      elements.append(get_next)
+
+    for i in range(4):
+      got = self.evaluate([element() for element in elements])
+      got.sort(key=lambda x: x[0])
+      expected = []
+      for j in range(100):
+        expected.append(range(i * 10000 + j * 100, i * 10000 + (j + 1) * 100))
+      self.assertAllEqual(got, expected)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate([element() for element in elements])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchSparse(self):
+
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+    dataset = dataset_ops.Dataset.range(10).apply(
+        batching.map_and_batch(_sparse, 5))
+
+    self.assertDatasetProduces(
+        dataset,
+        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)
+        ])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchFails(self):
+    """Test a dataset that maps a TF function across its input elements."""
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError, "oops"):
+      dataset = dataset_ops.Dataset.from_tensors(
+          array_ops.check_numerics(
+              constant_op.constant(1.0) / constant_op.constant(0.0), "oops"))
+      dataset = dataset.apply(batching.map_and_batch(lambda x: x, 14))
+      get_next = self.getNext(dataset, requires_initialization=True)
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchShapeMismatch(self):
+    """Test a dataset that maps a TF function across its input elements."""
+
+    def generator():
+      yield [1]
+      yield [2]
+      yield [3]
+      yield [[4, 5, 6]]
+
+    dataset = dataset_ops.Dataset.from_generator(
+        generator, output_types=dtypes.int32)
+    batch_size = 4
+    dataset = dataset.apply(batching.map_and_batch(lambda x: x, batch_size))
+    self.assertDatasetProduces(
+        dataset,
+        expected_error=(errors.InvalidArgumentError,
+                        "number of elements does not match"))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchImplicitDispose(self):
+    # Tests whether a map and batch dataset will be cleaned up correctly when
+    # the pipeline does not run it until exhaustion.
+    # The pipeline is TensorSliceDataset -> RepeatDataset(1000) ->
+    # MapAndBatchDataset(f=square_3, batch_size=100).
+    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).repeat(
+        1000).apply(batching.map_and_batch(_map_fn, batch_size=100))
+    dataset = dataset.prefetch(5)
+    get_next = self.getNext(dataset)
+    for _ in range(3):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(threshold=[0, 5, 10, 90, 95, 99]))
+  )
+  def testMapAndBatchMapError(self, threshold):
+
+    def raising_py_fn(i):
+      if i >= threshold:
+        raise StopIteration()
+      else:
+        return i
+
+    dataset = dataset_ops.Dataset.range(100).apply(
+        batching.map_and_batch(
+            lambda x: script_ops.py_func(raising_py_fn, [x], dtypes.int64),
+            batch_size=10))
+
+    get_next = self.getNext(dataset)
+    for i in range(threshold // 10):
+      self.assertAllEqual([i * 10 + j for j in range(10)],
+                          self.evaluate(get_next()))
+    for i in range(threshold // 10, 10):
+      with self.assertRaises(errors.InvalidArgumentError):
+        self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(element=False, dtype=dtypes.bool) +
+          combinations.combine(
+              element=-42,
+              dtype=[dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]) +
+          combinations.combine(element=42, dtype=[dtypes.uint8, dtypes.uint16])
+          + combinations.combine(
+              element=42.0,
+              dtype=[dtypes.float16, dtypes.float32, dtypes.float64]) +
+          combinations.combine(element=b"hello", dtype=[dtypes.string])))
+  def testMapAndBatchTypes(self, element, dtype):
+
+    def gen():
+      yield element
+
+    dataset = dataset_ops.Dataset.from_generator(gen, dtype).repeat(100).apply(
+        batching.map_and_batch(lambda x: x, batch_size=10))
+
+    get_next = self.getNext(dataset)
+    for _ in range(10):
+      self.assertAllEqual([element for _ in range(10)],
+                          self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuitIdentity(self):
+    map_fn = lambda x: x
+    dataset = self.structuredDataset(None).repeat().apply(
+        batching.map_and_batch(map_fn, batch_size=10))
+    get_next = self.getNext(dataset)
+    expected = map_fn(self.evaluate(self.structuredElement(None, shape=[10])))
+    self.assertAllEqual(expected, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuitReplicate(self):
+    map_fn = lambda x: (x, x)
+    dataset = self.structuredDataset(None).repeat().apply(
+        batching.map_and_batch(map_fn, batch_size=10))
+    get_next = self.getNext(dataset)
+    expected = map_fn(self.evaluate(self.structuredElement(None, shape=[10])))
+    self.assertAllEqual(expected, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuitSwap(self):
+    map_fn = lambda x, y: (y, x)
+    dataset = self.structuredDataset(
+        (None,
+         None)).repeat().apply(batching.map_and_batch(map_fn, batch_size=10))
+    get_next = self.getNext(dataset)
+    expected = map_fn(
+        *self.evaluate(self.structuredElement((None, None), shape=[10])))
+    self.assertAllEqual(expected, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuitProject(self):
+    map_fn = lambda x, y: x
+    dataset = self.structuredDataset(
+        (None,
+         None)).repeat().apply(batching.map_and_batch(map_fn, batch_size=10))
+    get_next = self.getNext(dataset)
+    expected = map_fn(
+        *self.evaluate(self.structuredElement((None, None), shape=[10])))
+    self.assertAllEqual(expected, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShortCircuitCapturedInput(self):
+    captured_t = variables.Variable(42)
+    dataset = self.structuredDataset(None).repeat().apply(
+        batching.map_and_batch(lambda x: captured_t, batch_size=10))
+    self.evaluate(variables.global_variables_initializer())
+    get_next = self.getNext(dataset, requires_initialization=True)
+    self.assertAllEqual([42] * 10, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchControlFlow(self):
+
+    def map_fn(x):
+      previous_control_flow_v2_value = control_flow_util.ENABLE_CONTROL_FLOW_V2
+      control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
+      return_value = cond.cond(x < 50, lambda: x + 1, lambda: x * x)
+      control_flow_util.ENABLE_CONTROL_FLOW_V2 = previous_control_flow_v2_value
+      return return_value
+
+    dataset = dataset_ops.Dataset.range(100).apply(
+        batching.map_and_batch(map_fn, batch_size=10))
+    get_next = self.getNext(dataset)
+    for i in range(10):
+      if i < 5:
+        self.assertAllEqual([i * 10 + j + 1 for j in range(10)],
+                            self.evaluate(get_next()))
+      else:
+        self.assertAllEqual(
+            [((i * 10) + j) * ((i * 10) + j) for j in range(10)],
+            self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testCheckpointLargeBatches(self):
+    if pywrap_sanitizers.is_tsan_enabled():
+      self.skipTest("Creating a large buffer causes OOM when using tsan.")
+    # Batches of size 512M
+    dataset = dataset_ops.Dataset.from_tensors(
+        array_ops.ones((64, 1024, 1024), dtype=dtypes.float32)).repeat()
+    dataset = dataset.map(lambda x: x+1, num_parallel_calls=5)
+    dataset = dataset.batch(2)
+    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()
+
+
+class MapAndBatchCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              drop_remainder=[True, False], symbolic_checkpoint=[True, False])))
+  def test(self, verify_fn, drop_remainder, symbolic_checkpoint):
+    range_size = 11
+    num_shards = 3
+    num_repeats = 2
+    batch_size = 5
+    num_parallel_calls = 7
+    total_outputs = (range_size // num_shards) * num_repeats
+    if drop_remainder:
+      num_outputs = total_outputs // batch_size
+    else:
+      num_outputs = int(math.ceil(total_outputs / batch_size))
+
+    def build_ds(range_start, drop_remainder=False, symbolic_checkpoint=False):
+
+      def _map_fn(x):
+        return math_ops.square(x)
+
+      dataset = dataset_ops.Dataset.range(
+          range_start, range_start + range_size)
+      dataset = dataset.shard(num_shards=num_shards, index=0)
+      dataset = dataset.repeat(num_repeats)
+      dataset = dataset.apply(
+          batching.map_and_batch(
+              map_func=_map_fn,
+              batch_size=batch_size,
+              num_parallel_calls=num_parallel_calls,
+              drop_remainder=drop_remainder))
+      options = options_lib.Options()
+      options.experimental_symbolic_checkpoint = symbolic_checkpoint
+      return dataset.with_options(options)
+
+    verify_fn(
+        self, lambda: build_ds(
+            10,
+            drop_remainder=drop_remainder,
+            symbolic_checkpoint=symbolic_checkpoint), num_outputs)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(drop_remainder=[True, False])))
+  def testNumParallelBatches(self, verify_fn, drop_remainder):
+    range_size = 11
+    num_shards = 3
+    num_repeats = 2
+    batch_size = 5
+    num_parallel_batches = 2
+    total_outputs = (range_size // num_shards) * num_repeats
+    if drop_remainder:
+      num_outputs = total_outputs // batch_size
+    else:
+      num_outputs = int(math.ceil(total_outputs / batch_size))
+
+    def build_ds(range_start, drop_remainder):
+
+      def _map_fn(x):
+        return math_ops.square(x)
+
+      return dataset_ops.Dataset.range(
+          range_start, range_start + range_size).shard(
+              num_shards=num_shards, index=0).repeat(num_repeats).apply(
+                  batching.map_and_batch(
+                      map_func=_map_fn,
+                      batch_size=batch_size,
+                      num_parallel_batches=num_parallel_batches,
+                      drop_remainder=drop_remainder))
+
+    verify_fn(self, lambda: build_ds(10, drop_remainder=drop_remainder),
+              num_outputs)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def testSparse(self, verify_fn):
+
+    def build_dataset():
+
+      def map_fn(i):
+        return sparse_tensor.SparseTensorValue(
+            indices=[[0]], values=(i * [1]), dense_shape=[1])
+
+      return dataset_ops.Dataset.range(10).apply(
+          batching.map_and_batch(map_fn, 5))
+
+    verify_fn(self, build_dataset, num_outputs=2)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,372 @@
+# 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 MapDefunOp."""
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import map_defun
+from tensorflow.python.data.kernel_tests import test_base
+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_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.platform import test
+
+
+# TODO(b/123903858): Add eager and V2 test coverage
+def _test_combinations():
+  return combinations.combine(tf_api_version=[1], mode=["graph"])
+
+
+class MapDefunTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(_test_combinations())
+  def testNoIntraOpLimit(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def simple_fn(x):
+      return x * 2 + 3
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(
+        simple_fn, [elems], [dtypes.int32], [(2,)],
+        max_intra_op_parallelism=0)[0]
+    expected = elems * 2 + 3
+    self.assertAllEqual(self.evaluate(r), self.evaluate(expected))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunSimple(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def simple_fn(x):
+      return x * 2 + 3
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(simple_fn, [elems], [dtypes.int32], [(2,)])[0]
+    expected = elems * 2 + 3
+    self.assertAllEqual(self.evaluate(r), self.evaluate(expected))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunMismatchedTypes(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def fn(x):
+      return math_ops.cast(x, dtypes.float64)
+
+    nums = [1, 2, 3, 4, 5, 6]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(fn, [elems], [dtypes.int32], [()])[0]
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(r)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunReduceDim(self):
+    # Tests where the output has a different rank from the input
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def fn(x):
+      return array_ops.gather(x, 0)
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(fn, [elems], [dtypes.int32], [()])[0]
+    expected = constant_op.constant([1, 3, 5])
+    self.assertAllEqual(self.evaluate(r), self.evaluate(expected))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunMultipleOutputs(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def fn(x):
+      return (x, math_ops.cast(x * 2 + 3, dtypes.float64))
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(fn, [elems], [dtypes.int32, dtypes.float64], [(2,),
+                                                                          (2,)])
+    expected = [elems, elems * 2 + 3]
+    self.assertAllEqual(self.evaluate(r), self.evaluate(expected))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunShapeInference(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def fn(x):
+      return x
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    result = map_defun.map_defun(fn, [elems], [dtypes.int32], [(2,)])[0]
+    self.assertEqual(result.get_shape(), (3, 2))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunPartialShapeInference(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def fn(x):
+      return x
+
+    elems = array_ops.placeholder(dtypes.int64, (None, 2))
+    result = map_defun.map_defun(fn, [elems], [dtypes.int32], [(2,)])
+    self.assertEqual(result[0].get_shape().as_list(), [None, 2])
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunRaisesErrorOnRuntimeShapeMismatch(self):
+
+    @def_function.function(input_signature=[
+        tensor_spec.TensorSpec(None, dtypes.int32),
+        tensor_spec.TensorSpec(None, dtypes.int32)
+    ])
+    def fn(x, y):
+      return x, y
+
+    elems1 = array_ops.placeholder(dtypes.int32)
+    elems2 = array_ops.placeholder(dtypes.int32)
+    result = map_defun.map_defun(fn, [elems1, elems2],
+                                 [dtypes.int32, dtypes.int32], [(), ()])
+    with self.cached_session() as sess:
+      with self.assertRaisesWithPredicateMatch(
+          errors.InvalidArgumentError,
+          "All inputs must have the same dimension 0."):
+        sess.run(result, feed_dict={elems1: [1, 2, 3, 4, 5], elems2: [1, 2, 3]})
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunRaisesDefunError(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def fn(x):
+      with ops.control_dependencies([check_ops.assert_equal(x, 0)]):
+        return array_ops.identity(x)
+
+    elems = constant_op.constant([0, 0, 0, 37, 0])
+    result = map_defun.map_defun(fn, [elems], [dtypes.int32], [()])
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(result)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunCancelledCorrectly(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([5], dtypes.int64)])
+    def defun(x):
+      # x has leading dimension 5, this will raise an error
+      return array_ops.gather(x, 10)
+
+    c = array_ops.tile(
+        array_ops.expand_dims(
+            constant_op.constant([1, 2, 3, 4, 5], dtype=dtypes.int64), 0),
+        [100, 1])
+    map_defun_op = map_defun.map_defun(defun, [c], [dtypes.int64], [()])[0]
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r"indices = 10 is not in \[0, 5\)"):
+      self.evaluate(map_defun_op)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithUnspecifiedOutputShape(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def simple_fn(x):
+      res = x * 2 + 3
+      return (res, res + 1, res + 2)
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(simple_fn, [elems],
+                            [dtypes.int32, dtypes.int32, dtypes.int32],
+                            [None, (None,), (2,)])
+    expected = elems * 2 + 3
+    self.assertAllEqual(self.evaluate(r[0]), self.evaluate(expected))
+    self.assertAllEqual(self.evaluate(r[1]), self.evaluate(expected + 1))
+    self.assertAllEqual(self.evaluate(r[2]), self.evaluate(expected + 2))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithDifferentOutputShapeEachRun(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
+    def simple_fn(x):
+      return x * 2 + 3
+
+    elems = array_ops.placeholder(dtypes.int32, name="data")
+    r = map_defun.map_defun(simple_fn, [elems], [dtypes.int32], [None])[0]
+    with session.Session() as sess:
+      self.assertAllEqual(sess.run(r, feed_dict={elems: [0]}), [3])
+      self.assertAllEqual(
+          sess.run(r, feed_dict={elems: [[0], [1]]}), [[3], [5]])
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithWrongOutputShape(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([2], dtypes.int32)])
+    def simple_fn(x):
+      return x * 2 + 3
+
+    nums = [[1, 2], [3, 4], [5, 6]]
+    elems = constant_op.constant(nums, dtype=dtypes.int32, name="data")
+    r = map_defun.map_defun(simple_fn, [elems], [dtypes.int32], [(1,)])[0]
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(r)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithInvalidInput(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
+    def simple_fn(x):
+      return x * 2
+
+    c = constant_op.constant(2)
+    with self.assertRaises(ValueError):
+      # Fails at graph construction time for inputs with known shapes.
+      r = map_defun.map_defun(simple_fn, [c], [dtypes.int32], [None])[0]
+    p = array_ops.placeholder(dtypes.int32)
+    r = map_defun.map_defun(simple_fn, [p], [dtypes.int32], [None])[0]
+    with session.Session() as sess:
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(r, feed_dict={p: 0})
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithParentCancellation(self):
+    # Checks that a cancellation of the parent graph is threaded through to
+    # MapDefunOp correctly.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def simple_fn(x):
+      del x
+      queue = data_flow_ops.FIFOQueue(10, dtypes.int32, ())
+      # Blocking
+      return queue.dequeue_many(5)
+
+    c = constant_op.constant([1, 2, 3, 4, 5])
+    map_defun_op = map_defun.map_defun(simple_fn, [c], [dtypes.int32], [()])[0]
+
+    with self.cached_session() as sess:
+      thread = self.checkedThread(
+          self.assert_op_cancelled, args=(map_defun_op,))
+      thread.start()
+      time.sleep(0.2)
+      sess.close()
+      thread.join()
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithCapturedInputs(self):
+    c = constant_op.constant(2)
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def fn(x):
+      return x + c
+
+    x = constant_op.constant([1, 2, 3, 4])
+    map_defun_op = map_defun.map_defun(fn, [x], [dtypes.int32], [()])[0]
+    expected = x + c
+    self.assertAllEqual(self.evaluate(expected), self.evaluate(map_defun_op))
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithVariantTensor(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.variant)])
+    def fn(x):
+      return x
+
+    st = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])
+
+    serialized = sparse_ops.serialize_sparse_v2(st, out_type=dtypes.variant)
+    serialized = array_ops_stack.stack([serialized, serialized])
+    map_defun_op = map_defun.map_defun(fn, [serialized], [dtypes.variant],
+                                       [None])[0]
+    deserialized = sparse_ops.deserialize_sparse(map_defun_op, dtypes.int32)
+    expected = sparse_tensor.SparseTensorValue(
+        indices=[[0, 0, 0], [0, 1, 2], [1, 0, 0], [1, 1, 2]],
+        values=[1, 2, 1, 2],
+        dense_shape=[2, 3, 4])
+    actual = self.evaluate(deserialized)
+    self.assertValuesEqual(expected, actual)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithVariantTensorAsCaptured(self):
+
+    st = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])
+    serialized = sparse_ops.serialize_sparse_v2(st, out_type=dtypes.variant)
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def fn(x):
+      del x
+      return serialized
+
+    x = constant_op.constant([0, 0])
+    map_defun_op = map_defun.map_defun(fn, [x], [dtypes.variant], [None])[0]
+    deserialized = sparse_ops.deserialize_sparse(map_defun_op, dtypes.int32)
+    expected = sparse_tensor.SparseTensorValue(
+        indices=[[0, 0, 0], [0, 1, 2], [1, 0, 0], [1, 1, 2]],
+        values=[1, 2, 1, 2],
+        dense_shape=[2, 3, 4])
+    actual = self.evaluate(deserialized)
+    self.assertValuesEqual(expected, actual)
+
+  @combinations.generate(_test_combinations())
+  def testMapDefunWithStrTensor(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
+    def fn(x):
+      return x
+
+    st = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])
+
+    serialized = sparse_ops.serialize_sparse_v2(st, out_type=dtypes.string)
+    serialized = array_ops_stack.stack([serialized, serialized])
+    map_defun_op = map_defun.map_defun(fn, [serialized], [dtypes.string],
+                                       [None])[0]
+    deserialized = sparse_ops.deserialize_sparse(map_defun_op, dtypes.int32)
+    expected = sparse_tensor.SparseTensorValue(
+        indices=[[0, 0, 0], [0, 1, 2], [1, 0, 0], [1, 1, 2]],
+        values=[1, 2, 1, 2],
+        dense_shape=[2, 3, 4])
+    actual = self.evaluate(deserialized)
+    self.assertValuesEqual(expected, actual)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,241 @@
+# 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 private `MatchingFilesDataset`."""
+import os
+import shutil
+import tempfile
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import matching_files
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.eager import context
+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 MatchingFilesDatasetTest(test_base.DatasetTestBase,
+                               parameterized.TestCase):
+
+  def setUp(self):
+    super(MatchingFilesDatasetTest, self).setUp()
+    self.tmp_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    shutil.rmtree(self.tmp_dir, ignore_errors=True)
+    super(MatchingFilesDatasetTest, self).tearDown()
+
+  def _touchTempFiles(self, filenames):
+    for filename in filenames:
+      open(os.path.join(self.tmp_dir, filename), 'a').close()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNonExistingDirectory(self):
+    """Test the MatchingFiles dataset with a non-existing directory."""
+
+    self.tmp_dir = os.path.join(self.tmp_dir, 'nonexistingdir')
+    dataset = matching_files.MatchingFilesDataset(
+        os.path.join(self.tmp_dir, '*'))
+    self.assertDatasetProduces(
+        dataset, expected_error=(errors.NotFoundError, ''))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDirectory(self):
+    """Test the MatchingFiles dataset with an empty directory."""
+
+    dataset = matching_files.MatchingFilesDataset(
+        os.path.join(self.tmp_dir, '*'))
+    self.assertDatasetProduces(
+        dataset, expected_error=(errors.NotFoundError, ''))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSimpleDirectory(self):
+    """Test the MatchingFiles dataset with a simple directory."""
+
+    filenames = ['a', 'b', 'c']
+    self._touchTempFiles(filenames)
+
+    dataset = matching_files.MatchingFilesDataset(
+        os.path.join(self.tmp_dir, '*'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(os.path.join(self.tmp_dir, filename))
+            for filename in filenames
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFileSuffixes(self):
+    """Test the MatchingFiles dataset using the suffixes of filename."""
+
+    filenames = ['a.txt', 'b.py', 'c.py', 'd.pyc']
+    self._touchTempFiles(filenames)
+
+    dataset = matching_files.MatchingFilesDataset(
+        os.path.join(self.tmp_dir, '*.py'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(os.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):
+    """Test the MatchingFiles dataset using the middles of filename."""
+
+    filenames = ['aa.txt', 'bb.py', 'bbc.pyc', 'cc.pyc']
+    self._touchTempFiles(filenames)
+
+    dataset = matching_files.MatchingFilesDataset(
+        os.path.join(self.tmp_dir, 'b*.py*'))
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[
+            compat.as_bytes(os.path.join(self.tmp_dir, filename))
+            for filename in filenames[1:3]
+        ],
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedDirectories(self):
+    """Test the MatchingFiles dataset with nested directories."""
+
+    filenames = []
+    width = 8
+    depth = 4
+    for i in range(width):
+      for j in range(depth):
+        new_base = os.path.join(self.tmp_dir, str(i),
+                                *[str(dir_name) for dir_name in range(j)])
+        os.makedirs(new_base)
+        child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
+        for f in child_files:
+          filename = os.path.join(new_base, f)
+          filenames.append(filename)
+          open(filename, 'w').close()
+
+    patterns = [
+        os.path.join(self.tmp_dir, os.path.join(*['**' for _ in range(depth)]),
+                     suffix) for suffix in ['*.txt', '*.log']
+    ]
+
+    dataset = matching_files.MatchingFilesDataset(patterns)
+    next_element = self.getNext(dataset)
+    expected_filenames = [
+        compat.as_bytes(filename)
+        for filename in filenames
+        if filename.endswith('.txt') or filename.endswith('.log')
+    ]
+    actual_filenames = []
+    while True:
+      try:
+        actual_filenames.append(compat.as_bytes(self.evaluate(next_element())))
+      except errors.OutOfRangeError:
+        break
+
+    self.assertCountEqual(expected_filenames, actual_filenames)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyPatterns(self):
+    """Test the MatchingFiles dataset with an empty patterns list."""
+    if context.executing_eagerly():
+      with self.assertRaises(errors.InvalidArgumentError):
+        matching_files.MatchingFilesDataset([])
+    else:
+      dataset = matching_files.MatchingFilesDataset([])
+      self.assertDatasetProduces(
+          dataset, expected_error=(errors.InvalidArgumentError, '')
+      )
+
+
+class MatchingFilesDatasetCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  def _build_iterator_graph(self, test_patterns):
+    return matching_files.MatchingFilesDataset(test_patterns)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    tmp_dir = tempfile.mkdtemp()
+    width = 16
+    depth = 8
+    for i in range(width):
+      for j in range(depth):
+        new_base = os.path.join(tmp_dir, str(i),
+                                *[str(dir_name) for dir_name in range(j)])
+        if not os.path.exists(new_base):
+          os.makedirs(new_base)
+        child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
+        for f in child_files:
+          filename = os.path.join(new_base, f)
+          open(filename, 'w').close()
+
+    patterns = [
+        os.path.join(tmp_dir, os.path.join(*['**'
+                                             for _ in range(depth)]), suffix)
+        for suffix in ['*.txt', '*.log']
+    ]
+
+    num_outputs = width * len(patterns)
+    verify_fn(self, lambda: self._build_iterator_graph(patterns), num_outputs)
+
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRestoreInvalidPatternIndex(self):
+    tmp_dir = tempfile.mkdtemp()
+
+    # Touch files
+    open(os.path.join(tmp_dir, 'a.txt'), 'w').close()
+    open(os.path.join(tmp_dir, 'b.txt'), 'w').close()
+
+    # Create an iterator with 2 patterns.
+    patterns1 = [
+        os.path.join(tmp_dir, 'a.txt'),
+        os.path.join(tmp_dir, 'b.txt'),
+    ]
+    ds1 = matching_files.MatchingFilesDataset(patterns1)
+    it1 = ds1.as_numpy_iterator()
+
+    # Consume elements of the dataset.
+    for _ in it1:
+      pass
+
+    # Save the iterator state
+    state = it1.save()
+
+    # Create a new dataset with only 1 pattern.
+    patterns2 = [os.path.join(tmp_dir, 'a.txt')]
+    ds2 = matching_files.MatchingFilesDataset(patterns2)
+    it2 = ds2.as_numpy_iterator()
+
+    # Restoring the state of it1 (current_pattern_index = 2) into it2
+    # (patterns size = 1) should raise an InvalidArgumentError.
+    with self.assertRaises(errors.InvalidArgumentError):
+      it2.restore(state)
+
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+if __name__ == '__main__':
+  test.main()

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

@@ -0,0 +1,61 @@
+# 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 private `_ModelDataset` transformation."""
+from absl.testing import parameterized
+
+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 map_op
+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 ModelDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAutotuneOption(self):
+    dataset = dataset_ops.Dataset.from_tensors(0)
+    dataset = dataset.map(lambda x: x).apply(
+        testing.assert_next(["Root"]))
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.autotune.enabled = True
+    dataset = dataset.with_options(options)
+    get_next = self.getNext(dataset)
+
+    self.assertEqual(0, self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelMapWithAutotune(self):
+    dataset = dataset_ops.Dataset.range(1000)
+    dataset = map_op._ParallelMapDataset(  # pylint: disable=protected-access
+        dataset,
+        lambda x: x + 1,
+        num_parallel_calls=1,
+        deterministic=True,
+        use_inter_op_parallelism=False)
+    dataset = dataset.map(
+        lambda x: x + 1, num_parallel_calls=dataset_ops.AUTOTUNE)
+    next_element = self.getNext(dataset)
+    self.evaluate(next_element())
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,56 @@
+# 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.experimental.non_serializable()`."""
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.platform import test
+
+
+class NonSerializableTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNonSerializable(self):
+    dataset = dataset_ops.Dataset.from_tensors(0)
+    dataset = dataset.apply(testing.assert_next(["FiniteSkip"]))
+    dataset = dataset.skip(0)  # Should not be removed by noop elimination
+    dataset = dataset.apply(testing.non_serializable())
+    dataset = dataset.apply(testing.assert_next(["MemoryCacheImpl"]))
+    dataset = dataset.skip(0)  # Should be removed by noop elimination
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.noop_elimination = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNonSerializableAsDirectInput(self):
+    """Tests that non-serializable dataset can be OptimizeDataset's input."""
+    dataset = dataset_ops.Dataset.from_tensors(0)
+    dataset = dataset.apply(testing.non_serializable())
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.noop_elimination = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/BUILD

@@ -0,0 +1,270 @@
+# Definitions are loaded separately so that copybara can pattern match (and modify) each definition.
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_py_strict_test(
+    name = "filter_fusion_test",
+    size = "medium",
+    srcs = ["filter_fusion_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "grappler_test",
+    size = "medium",
+    srcs = ["grappler_test.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_and_batch_fusion_test",
+    size = "small",
+    srcs = ["map_and_batch_fusion_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_and_filter_fusion_test",
+    size = "small",
+    srcs = ["map_and_filter_fusion_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_fusion_test",
+    size = "medium",
+    srcs = ["map_fusion_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "map_parallelization_test",
+    size = "small",
+    srcs = ["map_parallelization_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "filter_parallelization_test",
+    size = "medium",
+    srcs = ["filter_parallelization_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:map_fn",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "make_deterministic_test",
+    size = "small",
+    srcs = ["make_deterministic_test.py"],
+    deps = [
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/experimental/ops:interleave_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "noop_elimination_test",
+    size = "small",
+    srcs = ["noop_elimination_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:logging_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "optimization_test",
+    size = "medium",
+    srcs = ["optimization_test.py"],
+    shard_count = 2,
+    tags = ["nomsan"],  # Runs out of memory.
+    deps = [
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/experimental/ops:scan_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "shuffle_and_repeat_fusion_test",
+    size = "small",
+    srcs = ["shuffle_and_repeat_fusion_test.py"],
+    tags = [
+        "no_oss",
+        "no_pip",
+        "no_windows",
+    ],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "seq_interleave_prefetch_test",
+    size = "medium",
+    srcs = ["seq_interleave_prefetch_test.py"],
+    deps = [
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/experimental/ops:scan_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/data/experimental/kernel_tests/optimization/filter_fusion_test.py

@@ -0,0 +1,113 @@
+# 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 `FilterFusion` optimization."""
+import functools
+
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+  cases = []
+
+  take_all = lambda x: constant_op.constant(True)
+  is_zero = lambda x: math_ops.equal(x, 0)
+  greater = lambda x: math_ops.greater(x + 5, 0)
+  predicates = [take_all, is_zero, greater]
+  for i, x in enumerate(predicates):
+    for j, y in enumerate(predicates):
+      cases.append((lambda x: x, "Scalar{}{}".format(i, j), [x, y]))
+      for k, z in enumerate(predicates):
+        cases.append((lambda x: x, "Scalar{}{}{}".format(i, j, k), [x, y, z]))
+
+  take_all = lambda x, y: constant_op.constant(True)
+  is_zero = lambda x, y: math_ops.equal(x * math_ops.cast(y, dtypes.int64), 0)
+
+  cases.append((lambda x: (x, x), "Tuple1", [take_all, take_all]))
+  cases.append((lambda x: (x, 2), "Tuple2", [take_all, is_zero]))
+
+  def reduce_fn(x, y):
+    function, name, predicates = y
+    return x + combinations.combine(
+        function=function,
+        predicates=combinations.NamedObject(name, predicates))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class FilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_combinations()))
+  def testFilterFusion(self, function, predicates):
+    dataset = dataset_ops.Dataset.range(5).apply(
+        testing.assert_next(["Map", "Filter", "MemoryCacheImpl"])).map(function)
+    for predicate in predicates:
+      dataset = dataset.filter(predicate)
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.filter_fusion = True
+    dataset = dataset.with_options(options)
+    expected_output = []
+    for x in range(5):
+      r = function(x)
+      filtered = False
+      for predicate in predicates:
+        if isinstance(r, tuple):
+          b = predicate(*r)  # Pass tuple as multiple arguments.
+        else:
+          b = predicate(r)
+        if not self.evaluate(b):
+          filtered = True
+          break
+
+      if not filtered:
+        expected_output.append(r)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCapturedInputs(self):
+    a = constant_op.constant(3, dtype=dtypes.int64)
+    b = constant_op.constant(4, dtype=dtypes.int64)
+    some_tensor = math_ops.mul(a, b)
+
+    def predicate(y):
+      return math_ops.less(math_ops.cast(y, dtypes.int64), some_tensor)
+
+    # We currently do not support functions with captured inputs.
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.assert_next(["Filter", "Filter"
+                            ])).filter(predicate).filter(lambda x: True)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.filter_fusion = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=range(10))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/filter_parallelization_test.py

@@ -0,0 +1,291 @@
+# 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 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 dtypes
+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.ops import script_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 FilterParallelizationTest(test_base.DatasetTestBase,
+                                parameterized.TestCase):
+
+  def enableFilterParallelization(self, dataset):
+    options = options_lib.Options()
+    options.experimental_optimization.filter_parallelization = True
+    return dataset.with_options(options)
+
+  @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)
+      dataset = self.enableFilterParallelization(dataset)
+      dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+      # 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 = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    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 = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    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 = self.enableFilterParallelization(dataset)
+    # Calling `legacy_filter_fn` with `_predicate` makes the predicate passed to
+    # `Filter` stateful and therefore not parallelizable.
+    if repr(apply_filter) != "legacy_filter_fn":
+      dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    else:
+      dataset = dataset.apply(testing.assert_next(["Filter"]))
+    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 = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    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 = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    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 = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    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 testInputOutOfRange(self):
+
+    def py_fn(_):
+      raise StopIteration()
+
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    dataset = dataset.filter(
+        lambda x: script_ops.py_func(py_fn, [x], dtypes.bool, stateful=False))
+    get_next = self.getNext(dataset)
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(autotune=[False, True])))
+  def testAutotuneSetting(self, autotune):
+    dataset = dataset_ops.Dataset.range(4)
+    options = options_lib.Options()
+    options.experimental_optimization.filter_parallelization = True
+    options.autotune.enabled = autotune
+    dataset = dataset.with_options(options)
+    if autotune:
+      dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    else:
+      dataset = dataset.apply(testing.assert_next(["Filter"]))
+    dataset = dataset.filter(
+        lambda x: math_ops.not_equal(math_ops.mod(x, 3), 2))
+    self.assertDatasetProduces(dataset, expected_output=[0, 1, 3])
+
+
+class FilterCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                           parameterized.TestCase):
+
+  def enableFilterParallelization(self, dataset):
+    options = options_lib.Options()
+    options.experimental_optimization.filter_parallelization = True
+    return dataset.with_options(options)
+
+  def _build_filter_range_graph(self, div):
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    return dataset.filter(
+        lambda x: math_ops.not_equal(math_ops.mod(x, div), 2))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    div = 3
+    num_outputs = sum(x % 3 != 2 for x in range(100))
+    verify_fn(self, lambda: self._build_filter_range_graph(div), num_outputs)
+
+  def _build_filter_dict_graph(self):
+    # pylint: disable=g-long-lambda
+    dataset = dataset_ops.Dataset.range(10).map(lambda x: {
+        "foo": x * 2,
+        "bar": x**2
+    })
+    dataset = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    return dataset.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))
+    # pylint: disable=unnecessary-lambda
+    verify_fn(self, lambda: self._build_filter_dict_graph(), num_outputs)
+
+  def _build_sparse_filter(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)
+
+    dataset = dataset_ops.Dataset.range(10).map(_map_fn)
+    dataset = self.enableFilterParallelization(dataset)
+    dataset = dataset.apply(testing.assert_next(["ParallelFilter"]))
+    return dataset.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):
+    # pylint: disable=unnecessary-lambda
+    verify_fn(self, lambda: self._build_sparse_filter(), num_outputs=5)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/grappler_test.py

@@ -0,0 +1,77 @@
+# 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 the generic Grappler optimizations used within tf.data."""
+from absl.testing import parameterized
+
+from tensorflow.core.example import example_pb2
+from tensorflow.core.example import feature_pb2
+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 test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.platform import test
+
+
+class GrapplerTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConstantFoldingVarLenFeature(self):
+    example = example_pb2.Example(features=feature_pb2.Features(feature={}))
+    dataset = dataset_ops.Dataset.from_tensors(example.SerializeToString())
+
+    def parse_fn(serialized):
+      features = {"x": parsing_ops.VarLenFeature(dtypes.int64)}
+      parsed = parsing_ops.parse_single_example(serialized, features)
+      parsed = parsed["x"].values
+
+      size = array_ops.size(parsed)
+      value = math_ops.cast(parsed, dtypes.bool)
+      return cond.cond(size > 0,
+                       lambda: array_ops.reshape(value, []),
+                       lambda: array_ops.zeros([], dtypes.bool))
+
+    dataset = dataset.map(parse_fn)
+
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLayoutOptimizationConv2D(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    # Compute convolution with input and filter of [1, 1, 1, 1] shape.
+    # Verify that Grappler doesn't transpose Conv2D data format to NCHW.
+    dataset = dataset_ops.Dataset.from_tensors((1, 1))
+
+    def map_function(x, y):
+      i = math_ops.cast(x, dtypes.float32)
+      i = array_ops.reshape(i, [1, 1, 1, 1])
+      f = math_ops.cast(y, dtypes.float32)
+      f = array_ops.reshape(f, [1, 1, 1, 1])
+      c = nn_ops.conv2d(i, f, strides=[1, 1, 1, 1], padding="VALID")
+      return array_ops.reshape(c, ())
+
+    dataset = dataset.map(map_function)
+    self.assertDatasetProduces(dataset, expected_output=[1])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/make_deterministic_test.py

@@ -0,0 +1,389 @@
+# Copyright 2021 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 `MakeDeterministic` optimization."""
+import os
+import time
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import interleave_ops
+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.data.ops import readers as reader_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 dtypes
+from tensorflow.python.framework import random_seed
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class MakeDeterministicTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _set_seed(self):
+    # Set the seed, since in graph mode some non-random dataset ops call
+    # tf.compat.v1.get_seed to copy the seed to a Defun. Calling get_seed raises
+    # an error with determinism if no seed is set.
+    # TODO(reedwm): Ensure such dataset ops do not raise an error when no seed
+    # is set.
+    random_seed.set_random_seed(1)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              use_function=[False, True], use_legacy_interleave=[False, True])))
+  def test_stateful_ops_interleave(self, use_function, use_legacy_interleave):
+    with test_util.deterministic_ops():
+
+      v = variables.Variable(0.)
+
+      def map_fn(x):
+        v.assign_add(1.)
+        return (x, v.read_value())
+
+      def interleave_fn(x):
+        del x
+        return dataset_ops.Dataset.range(2).map(map_fn)
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+        interleave_fn = def_function.function(interleave_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      if use_legacy_interleave:
+        dataset = dataset.apply(
+            interleave_ops.parallel_interleave(interleave_fn, cycle_length=5))
+      else:
+        dataset = dataset.interleave(
+            interleave_fn, cycle_length=5, num_parallel_calls=3)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = list(zip([0] * 5 + [1] * 5, range(1, 11)))
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=expected_output,
+          requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              use_function=[False, True])))
+  def test_stateful_ops_map(self, use_function):
+    with test_util.deterministic_ops():
+
+      v = variables.Variable(0.)
+
+      def map_fn(x):
+        v.assign_add(1.)
+        return (x, v.read_value())
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.map(map_fn, num_parallel_calls=5)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = list(zip(range(0, 5), range(1, 6)))
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=expected_output,
+          requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_stateful_ops_map_with_random_ops(self):
+    with test_util.deterministic_ops():
+
+      def map_fn(x):
+        return x + random_ops.random_uniform(
+            (), 0, 2, dtype=dtypes.int64, seed=1)
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.apply(testing.assert_next(["Map", "ParallelMap"]))
+      dataset = dataset.map(map_fn, num_parallel_calls=5)
+      get_next = self.getNext(dataset, requires_initialization=True)
+      for i in range(5):
+        self.assertIn(self.evaluate(get_next()), [i, i + 1])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(use_function=[False, True])))
+  def test_stateful_ops_map_ignore_input(self, use_function):
+    with test_util.deterministic_ops():
+
+      v = variables.Variable(0.)
+
+      def map_fn(x):
+        del x
+        v.assign_add(1.)
+        return math_ops.constant_op.constant(1.)
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.map(map_fn, num_parallel_calls=5)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = [1.] * 5
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=expected_output,
+          requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(use_function=[False, True])))
+  def test_stateful_ops_batch(self, use_function):
+    with test_util.deterministic_ops():
+
+      v = variables.Variable(0.)
+
+      def map_fn(x):
+        return (x, v.read_value())
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.map(map_fn)
+      dataset = dataset.apply(testing.assert_next(["Batch"]))
+      dataset = dataset.batch(2, num_parallel_calls=2)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = [
+          (np.array([0, 1]), np.array([0, 0])),
+          (np.array([2, 3]), np.array([0, 0])),
+          (np.array([4]), np.array([0])),
+      ]
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=expected_output,
+          requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              use_function=[False, True],
+              use_legacy_map_and_batch=[False, True])))
+  def test_stateful_ops_map_and_batch(self, use_function,
+                                      use_legacy_map_and_batch):
+    with test_util.deterministic_ops():
+
+      v = variables.Variable(0.)
+
+      def map_fn(x):
+        v.assign_add(1.)
+        return (x, v.read_value())
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      if use_legacy_map_and_batch:
+        dataset = dataset.apply(batching.map_and_batch(map_fn, 2,
+                                                       num_parallel_calls=5))
+      else:
+        dataset = dataset.map(map_fn, num_parallel_calls=5)
+        dataset = dataset.batch(2)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = [
+          (np.array([0, 1]), np.array([1, 2])),
+          (np.array([2, 3]), np.array([3, 4])),
+          (np.array([4]), np.array([5])),
+      ]
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=expected_output,
+          requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              use_function=[False, True], use_legacy_interleave=[False, True])))
+  def test_no_stateful_ops_interleave(self, use_function,
+                                      use_legacy_interleave):
+    self._set_seed()
+    with test_util.deterministic_ops():
+
+      def interleave_fn(x):
+        del x
+        return dataset_ops.Dataset.range(2)
+
+      if use_function:
+        interleave_fn = def_function.function(interleave_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      if use_legacy_interleave:
+        dataset = dataset.apply(
+            testing.assert_next(["LegacyParallelInterleaveV2"]))
+        dataset = dataset.apply(
+            interleave_ops.parallel_interleave(interleave_fn, cycle_length=5))
+      else:
+        dataset = dataset.apply(testing.assert_next(["ParallelInterleave"]))
+        dataset = dataset.interleave(
+            interleave_fn, cycle_length=5, num_parallel_calls=3)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertDatasetProduces(dataset, expected_output=[0] * 5 + [1] * 5)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(use_function=[False, True])))
+  def test_no_stateful_ops_map(self, use_function):
+    self._set_seed()
+    with test_util.deterministic_ops():
+      def map_fn(x):
+        return x + 1
+
+      if use_function:
+        map_fn = def_function.function(map_fn)
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.apply(testing.assert_next(["ParallelMap"]))
+      dataset = dataset.map(map_fn, num_parallel_calls=5)
+      self.evaluate(variables.global_variables_initializer())
+      expected_output = range(1, 6)
+      self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              use_function=[False, True], use_control_flow=[False, True])))
+  def test_text_line_dataset(self, use_function, use_control_flow):
+    self._set_seed()
+    with test_util.deterministic_ops():
+
+      def write_nums_to_file(filename, numbers):
+        path = os.path.join(self.get_temp_dir(), filename)
+        with open(path, "w") as f:
+          f.write("\n".join(str(n) for n in numbers))
+        return path
+
+      f1 = write_nums_to_file("f1", (1, 2, 3))
+      f2 = write_nums_to_file("f2", (4, 5, 6))
+      f3 = write_nums_to_file("f3", (7, 8, 9))
+
+      if use_control_flow:
+        def interleave_fn(filename):
+          # Test function that uses control flow. The True branch is never taken
+          concat = string_ops.string_join([filename, "abc"])
+          return cond.cond(
+              math_ops.equal(filename, "abc"),
+              lambda: reader_ops.TextLineDataset(concat),
+              lambda: reader_ops.TextLineDataset(filename))
+      else:
+        def interleave_fn(filename):
+          return reader_ops.TextLineDataset(filename)
+
+      if use_function:
+        interleave_fn = def_function.function(interleave_fn)
+
+      dataset = dataset_ops.Dataset.from_tensor_slices([f1, f2, f3])
+      dataset = dataset.apply(testing.assert_next(["ParallelInterleave"]))
+      dataset = dataset.interleave(
+          interleave_fn, cycle_length=3, num_parallel_calls=3)
+
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=["1", "4", "7", "2", "5", "8", "3", "6", "9"])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              local_determinism=[None, True, False],
+              global_determinism=[True, False])))
+  def test_deterministic_attribute(self, local_determinism, global_determinism):
+    self._set_seed()
+    with test_util.deterministic_ops():
+
+      def sleep(x):
+        time.sleep(0.1)
+        return x
+
+      def map_function(x):
+        if math_ops.equal(x, 0):
+          return script_ops.py_func(sleep, [x], x.dtype, stateful=False)
+        else:
+          return x
+
+      dataset = dataset_ops.Dataset.range(100)
+      dataset = dataset.map(
+          map_function, num_parallel_calls=2, deterministic=local_determinism)
+      opts = options_lib.Options()
+      opts.deterministic = global_determinism
+      dataset = dataset.with_options(opts)
+
+      self.assertDatasetProduces(dataset, expected_output=range(100))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_rewrite_prefetch(self):
+    with test_util.deterministic_ops():
+      v = variables.Variable(-1, dtype=dtypes.int64)
+
+      def map_fn(x):
+        v.assign(x)
+        return x
+
+      dataset = dataset_ops.Dataset.range(5)
+      dataset = dataset.map(map_fn)
+      dataset = dataset.prefetch(5)
+      self.evaluate(variables.global_variables_initializer())
+      get_next = self.getNext(dataset, requires_initialization=True)
+      self.assertEqual(self.evaluate(v), -1)
+      self.assertEqual(self.evaluate(get_next()), 0)
+      time.sleep(0.01)
+      self.assertEqual(self.evaluate(v), 0)
+      self.assertEqual(self.evaluate(get_next()), 1)
+      time.sleep(0.01)
+      self.assertEqual(self.evaluate(v), 1)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_no_determinism(self):
+    config.disable_op_determinism()
+    v = variables.Variable(0.)
+
+    def interleave_fn(x):
+      del x
+      v.assign(1.)
+      return dataset_ops.Dataset.range(2)
+
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.apply(testing.assert_next(["ParallelInterleave"]))
+    dataset = dataset.interleave(
+        interleave_fn, cycle_length=5, num_parallel_calls=3)
+    self.evaluate(variables.global_variables_initializer())
+    expected_output = [0] * 5 + [1] * 5
+    self.assertDatasetProduces(
+        dataset, expected_output=expected_output, requires_initialization=True)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/map_and_batch_fusion_test.py

@@ -0,0 +1,42 @@
+# 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 `MapAndBatchFusion` optimization."""
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.platform import test
+
+
+class MapAndBatchFusionTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapAndBatchFusion(self):
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.assert_next(
+            ["MapAndBatch"])).map(lambda x: x * x).batch(10)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_and_batch_fusion = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset, expected_output=[[x * x for x in range(10)]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/map_and_filter_fusion_test.py

@@ -0,0 +1,122 @@
+# 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 `MapAndFilterFusion` optimization."""
+import functools
+
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+  cases = []
+
+  identity = lambda x: x
+  increment = lambda x: x + 1
+  minus_five = lambda x: x - 5
+
+  def increment_and_square(x):
+    y = x + 1
+    return y * y
+
+  functions = [identity, increment, minus_five, increment_and_square]
+
+  take_all = lambda x: constant_op.constant(True)
+  is_zero = lambda x: math_ops.equal(x, 0)
+  is_odd = lambda x: math_ops.equal(x % 2, 0)
+  greater = lambda x: math_ops.greater(x + 5, 0)
+  predicates = [take_all, is_zero, is_odd, greater]
+
+  for i, function in enumerate(functions):
+    for j, predicate in enumerate(predicates):
+      cases.append((function, "Scalar{}{}".format(i, j), predicate))
+
+  replicate = lambda x: (x, x)
+  with_two = lambda x: (x, 2)
+  functions = [replicate, with_two]
+  take_all = lambda x, y: constant_op.constant(True)
+  is_zero = lambda x, y: math_ops.equal(x * math_ops.cast(y, dtypes.int64), 0)
+  predicates = [take_all, is_zero]
+
+  for i, function in enumerate(functions):
+    for j, predicate in enumerate(predicates):
+      cases.append((function, "Tuple{}{}".format(i, j), predicate))
+
+  def reduce_fn(x, y):
+    function, name, predicate = y
+    return x + combinations.combine(
+        function=function,
+        predicate=combinations.NamedObject(name, predicate))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class MapAndFilterFusionTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _testDataset(self, dataset, function, predicate):
+    expected_output = []
+    for x in range(10):
+      r = function(x)
+      if isinstance(r, tuple):
+        b = predicate(*r)  # Pass tuple as multiple arguments.
+      else:
+        b = predicate(r)
+      if self.evaluate(b):
+        expected_output.append(r)
+    self.assertDatasetProduces(dataset, expected_output=expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_combinations()))
+  def testMapAndFilterFusion(self, function, predicate):
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.assert_next(["Map", "Filter",
+                             "Map"])).map(function).filter(predicate)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_and_filter_fusion = True
+    dataset = dataset.with_options(options)
+    self._testDataset(dataset, function, predicate)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCapturedInputs(self):
+    a = constant_op.constant(3, dtype=dtypes.int64)
+    b = constant_op.constant(4, dtype=dtypes.int64)
+    some_tensor = math_ops.mul(a, b)
+    function = lambda x: x * x
+
+    def predicate(y):
+      return math_ops.less(math_ops.cast(y, dtypes.int64), some_tensor)
+
+    # We currently do not support functions with captured inputs.
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.assert_next(["Map", "Filter"])).map(function).filter(predicate)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_and_filter_fusion = True
+    dataset = dataset.with_options(options)
+    self._testDataset(dataset, function, predicate)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/map_fusion_test.py

@@ -0,0 +1,266 @@
+# 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 `MapFusion` optimization."""
+import functools
+
+from absl.testing import parameterized
+
+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.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 check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+  cases = []
+
+  identity = lambda x: x
+  increment = lambda x: x + 1
+
+  def increment_and_square(x):
+    y = x + 1
+    return y * y
+
+  functions = [identity, increment, increment_and_square]
+
+  for i, x in enumerate(functions):
+    for j, y in enumerate(functions):
+      cases.append(("Scalar{}{}".format(i, j), [x, y]))
+      for k, z in enumerate(functions):
+        cases.append(("Scalar{}{}{}".format(i, j, k), [x, y, z]))
+
+  with_42 = lambda x: (x, 42)
+  swap = lambda x, y: (y, x)
+
+  cases.append(("Tuple1", [with_42, swap]))
+  cases.append(("Tuple2", [with_42, swap, swap]))
+
+  def reduce_fn(x, y):
+    name, functions = y
+    return x + combinations.combine(
+        functions=combinations.NamedObject(name, functions)
+    )
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class MapFusionTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          _test_combinations(),
+          combinations.combine(
+              num_parallel_calls=[None, 2, dataset_ops.AUTOTUNE]
+          ),
+          combinations.combine(deterministic=[None, True, False]),
+      )
+  )
+  def testMapFusion(self, functions, num_parallel_calls, deterministic):
+    dataset = dataset_ops.Dataset.range(5)
+    if num_parallel_calls is None:
+      dataset = dataset.apply(testing.assert_next(["Map", "MemoryCacheImpl"]))
+    elif num_parallel_calls in [dataset_ops.AUTOTUNE]:
+      # TODO(b/148614504): Support fusion of parallel maps with
+      # non-AUTOTUNE value.
+      dataset = dataset.apply(
+          testing.assert_next(["ParallelMap", "MemoryCacheImpl"])
+      )
+    else:
+      dataset = dataset.apply(
+          testing.assert_next(["ParallelMap", "ParallelMap"])
+      )
+
+    for function in functions:
+      dataset = dataset.map(
+          function,
+          num_parallel_calls=num_parallel_calls,
+          deterministic=deterministic,
+      )
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+    expected_output = []
+    for x in range(5):
+      r = x
+      for function in functions:
+        if isinstance(r, tuple):
+          r = function(*r)  # Pass tuple as multiple arguments.
+        else:
+          r = function(r)
+      expected_output.append(r)
+
+    nondeterministic_ordering = (
+        num_parallel_calls is not None and deterministic is False  # pylint: disable=g-bool-id-comparison
+    )
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=expected_output,
+        assert_items_equal=nondeterministic_ordering,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapFusionLongMapChain(self):
+    n = 5
+    dataset = dataset_ops.Dataset.range(n)
+    dataset = dataset.apply(
+        testing.assert_next(["ParallelMap", "MemoryCacheImpl"])
+    )
+
+    k = 50
+    for _ in range(k):
+      dataset = dataset.map(
+          lambda x: 2 * x,
+          num_parallel_calls=dataset_ops.AUTOTUNE,
+      )
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[x * 2**k for x in range(n)],
+        assert_items_equal=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMixedMapAndParallelMap(self):
+    n = 5
+    dataset = dataset_ops.Dataset.range(n)
+    dataset = dataset.apply(
+        testing.assert_next(["ParallelMap", "MemoryCacheImpl"])
+    )
+
+    # These 3 maps should be merged because the non-parallel Map should be
+    # parallelized by `map_parallelization` before `map_fusion` is applied.
+    dataset = dataset.map(lambda x: 2 * x,
+                          num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.map(lambda x: 2 * x)
+    dataset = dataset.map(lambda x: 2 * x,
+                          num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[x * 2**3 for x in range(n)],
+        assert_items_equal=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testControlInputs(self):
+    def f(x):
+      with ops.control_dependencies([check_ops.assert_type(x, dtypes.int64)]):
+        return 2 * x
+
+    n = 5
+    dataset = dataset_ops.Dataset.range(n)
+    dataset = dataset.apply(
+        testing.assert_next(["ParallelMap", "MemoryCacheImpl"])
+    )
+    dataset = dataset.map(f, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.map(f, num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[x * 4 for x in range(n)],
+        assert_items_equal=True,
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStatefulness(self):
+    def f(x):
+      return control_flow_ops.with_dependencies(
+          [check_ops.assert_negative(x)], x
+      )
+
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.apply(
+        testing.assert_next(["ParallelMap", "MemoryCacheImpl"])
+    )
+    dataset = dataset.map(lambda x: x, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.map(f, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.map(lambda x: x, num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    dataset = dataset.cache()
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "assertion failed"
+    ):
+      self.evaluate(self.getNext(dataset)())
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_parallel_calls=[None, 2, dataset_ops.AUTOTUNE]
+          ),
+      )
+  )
+  def testCapturedInputs(self, num_parallel_calls):
+    a = constant_op.constant(3, dtype=dtypes.int64)
+    b = constant_op.constant(4, dtype=dtypes.int64)
+    some_tensor = math_ops.mul(a, b)
+
+    dataset = dataset_ops.Dataset.range(1)
+    # We currently do not support functions with captured inputs.
+    if num_parallel_calls in [2, dataset_ops.AUTOTUNE]:
+      dataset = dataset.apply(
+          testing.assert_next(["ParallelMap", "ParallelMap"])
+      )
+    else:
+      dataset = dataset.apply(testing.assert_next(["Map", "Map"]))
+
+    dataset = dataset.map(
+        lambda x: some_tensor, num_parallel_calls=num_parallel_calls
+    ).map(lambda x: x, num_parallel_calls=num_parallel_calls)
+
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_fusion = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[some_tensor])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/map_parallelization_test.py

@@ -0,0 +1,169 @@
+# 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 `MapParallelization` optimization."""
+import functools
+
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+  def assert_greater(x):
+    assert_op = control_flow_assert.Assert(math_ops.greater(x, -1), [x])
+    with ops.control_dependencies([assert_op]):
+      return x
+
+  cases = [
+      ("Identity", lambda x: x, True),
+      ("Increment", lambda x: x + 1, True),
+      ("AssertGreater", assert_greater, True),
+  ]
+
+  def reduce_fn(x, y):
+    name, function, should_optimize = y
+    return x + combinations.combine(
+        function=combinations.NamedObject(name, function),
+        should_optimize=should_optimize)
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class MapParallelizationTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_combinations()))
+  def testMapParallelization(self, function, should_optimize):
+    next_nodes = ["ParallelMap"] if should_optimize else ["Map"]
+    dataset = dataset_ops.Dataset.range(5).apply(
+        testing.assert_next(next_nodes)).map(function)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset, expected_output=[function(x) for x in range(5)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoMapParallelizationWhenSynchronous(self):
+    dataset = (
+        dataset_ops.Dataset.range(5)
+        .apply(testing.assert_next(["Map"]))
+        .map(lambda x: x + 1, synchronous=True)
+    )
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset, expected_output=[x + 1 for x in range(5)]
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCapturedConstant(self):
+    captured_t = constant_op.constant(42, dtype=dtypes.int64)
+    def fn(x):
+      return x + captured_t
+    dataset = dataset_ops.Dataset.range(5).apply(
+        testing.assert_next(["ParallelMap"])).map(fn)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(
+        dataset, expected_output=[x + 42 for x in range(5)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCapturedVariable(self):
+    captured_t = variables.Variable(42, dtype=dtypes.int64)
+    def fn(x):
+      return x + captured_t
+    dataset = dataset_ops.Dataset.range(5).apply(
+        testing.assert_next(["Map"])).map(fn)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+    self.evaluate(variables.global_variables_initializer())
+    self.assertDatasetProduces(
+        dataset,
+        expected_output=[x + 42 for x in range(5)],
+        requires_initialization=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(apply_autotune=[None, True, False])))
+  def testAutotuneOption(self, apply_autotune):
+    next_nodes = ["ParallelMap"] if (apply_autotune is not False) else ["Map"]  # pylint: disable=g-bool-id-comparison
+    dataset = dataset_ops.Dataset.range(4).apply(
+        testing.assert_next(next_nodes)).map(lambda x: x + 2)
+
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    if apply_autotune is not None:
+      options.autotune.enabled = apply_autotune
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[2, 3, 4, 5])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoParallelizationInsideInterleave(self):
+
+    def func(i):
+      ds = dataset_ops.Dataset.range(i).apply(testing.assert_next(
+          ["Map"])).map(lambda x: x + 1)
+      return ds
+
+    dataset = dataset_ops.Dataset.range(1, 4).interleave(
+        map_func=func, cycle_length=2, block_length=2)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(dataset, expected_output=[1, 1, 2, 1, 2, 3])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoParallelizationInsideFlatMap(self):
+
+    def func(i):
+      ds = dataset_ops.Dataset.range(i).apply(testing.assert_next(
+          ["Map"])).map(lambda x: x + 1)
+      return ds
+
+    dataset = dataset_ops.Dataset.range(1, 4).flat_map(map_func=func)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_parallelization = True
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(dataset, expected_output=[1, 1, 2, 1, 2, 3])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/noop_elimination_test.py

@@ -0,0 +1,131 @@
+# 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 `NoopElimination` optimization."""
+import functools
+
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.platform import test
+
+
+def _test_combinations():
+
+  def make_range():
+    return dataset_ops.Dataset.range(10)
+
+  def fn_with_side_effect(arg):
+    logging_ops.print_v2(arg)
+    return arg
+
+  # Test case for map function with capture args
+  def apply_map_with_capture(ds):
+    const = constant_op.constant(-1, dtype=dtypes.int64)
+    return ds.map(lambda x: (x, const))
+
+  # Test case for map functions with multiple components
+  def apply_map_with_multiple_components(ds):
+    ds = ds.map(lambda x: (x, x), num_parallel_calls=2)  # Not eliminated
+    return ds.map(lambda x, y: (x, y))  # Eliminated
+
+  parallel_map_name = "ParallelMap"
+
+  cases = [
+      ("Skip0", lambda ds: ds.skip(0), None),
+      ("SkipN", lambda ds: ds.skip(5), "FiniteSkip"),
+      ("Repeat1", lambda ds: ds.repeat(1), None),
+      ("RepeatN", lambda ds: ds.repeat(10), "FiniteRepeat[0]"),
+      ("Prefetch0", lambda ds: ds.prefetch(0), "Prefetch"),
+      ("PrefetchN", lambda ds: ds.prefetch(1), "Prefetch"),
+      ("Take-1", lambda ds: ds.take(-1), None),
+      ("TakeN", lambda ds: ds.take(2), "FiniteTake"),
+      ("MapIdentity", lambda ds: ds.map(lambda x: x), None),
+      ("MapNonIdentity", lambda ds: ds.map(lambda x: x * 2), "Map"),
+      ("MapWithSideEffect", lambda ds: ds.map(fn_with_side_effect), "Map"),
+      ("MapWithCapture", apply_map_with_capture, "Map"),
+      (
+          "MapWithMultipleComponents",
+          apply_map_with_multiple_components,
+          parallel_map_name,
+      ),
+      ("MapRestructure", lambda ds: ds.map(lambda x: {"value": x}), ""),
+      (
+          "PMapIdentity",
+          lambda ds: ds.map(lambda x: x, num_parallel_calls=2),
+          None,
+      ),
+      (
+          "PMapNonIdentity",
+          lambda ds: ds.map(lambda x: x * 2, num_parallel_calls=2),
+          parallel_map_name,
+      ),
+      ("Shard1", lambda ds: ds.shard(1, 0), None),
+      ("ShardN", lambda ds: ds.shard(2, 0), "Shard"),
+  ]
+
+  def reduce_fn(result, case):
+    name, transformation, expected = case
+    return result + combinations.combine(
+        init_dataset_fn=make_range,
+        transformation=combinations.NamedObject(name, transformation),
+        expected_name=expected)
+
+  test_combinations = functools.reduce(reduce_fn, cases, [])
+
+  return test_combinations
+
+
+class NoopEliminationTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_combinations()))
+  def testNoopElimination(self, init_dataset_fn, transformation, expected_name):
+    """Runs a noop elimination test case.
+
+    Args:
+      init_dataset_fn: Function to create the initial dataset
+      transformation: Transformation to apply
+      expected_name: Name of the transformation if it is not eliminated
+    """
+    dataset = init_dataset_fn()
+
+    if expected_name:
+      dataset = dataset.apply(
+          testing.assert_next([expected_name, "FiniteTake"]))
+    else:
+      dataset = dataset.apply(testing.assert_next(["FiniteTake"]))
+
+    dataset = dataset.apply(transformation)
+    dataset = dataset.take(1)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.noop_elimination = True
+    dataset = dataset.with_options(options)
+
+    # Run the first iteration for the side effect of checking the assertion.
+    get_next = self.getNext(dataset)
+    self.evaluate(get_next())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/optimization_test.py

@@ -0,0 +1,308 @@
+# 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 static tf.data optimizations."""
+import functools
+import time
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import batching
+from tensorflow.python.data.experimental.ops import grouping
+from tensorflow.python.data.experimental.ops import scan_ops
+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.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 array_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def _captured_refvar_test_combinations():
+
+  def make_map_dataset(var):
+    return dataset_ops.Dataset.from_tensors(0).map(lambda x: x + var)
+
+  def make_flat_map_dataset(var):
+    return dataset_ops.Dataset.from_tensors(
+        0).flat_map(lambda _: dataset_ops.Dataset.from_tensors(var))
+
+  def make_filter_dataset(var):
+    return dataset_ops.Dataset.from_tensors(0).filter(lambda x: x < var)
+
+  def make_map_and_batch_dataset(var):
+
+    def map_fn(x):
+      return x + var
+
+    return dataset_ops.Dataset.from_tensors(0).apply(
+        batching.map_and_batch(map_fn, 1))
+
+  def make_group_by_reducer_dataset(var):
+    reducer = grouping.Reducer(
+        init_func=lambda _: 0,
+        reduce_func=lambda x, y: x,
+        finalize_func=lambda _: var)
+    return dataset_ops.Dataset.range(5).apply(
+        grouping.group_by_reducer(lambda x: x % 2, reducer))
+
+  def make_group_by_window_dataset(var):
+
+    def reduce_fn(key, bucket):
+      del key, bucket
+      return dataset_ops.Dataset.from_tensors(var)
+
+    return dataset_ops.Dataset.from_tensors(0).repeat(10).apply(
+        grouping.group_by_window(lambda _: 0, reduce_fn, 10))
+
+  def make_scan_dataset(var):
+    return dataset_ops.Dataset.from_tensors(0).apply(
+        scan_ops.scan(
+            0, lambda old_state, elem: (old_state + 1, elem + old_state + var)))
+
+  cases = [
+      # Core datasets
+      ("Map", make_map_dataset),
+      ("FlatMap", make_flat_map_dataset),
+      ("Filter", make_filter_dataset),
+      # Experimental datasets
+      ("MapAndBatch", make_map_and_batch_dataset),
+      ("GroupByReducer", make_group_by_reducer_dataset),
+      ("GroupByWindow", make_group_by_window_dataset),
+      ("Scan", make_scan_dataset)
+  ]
+
+  def reduce_fn(x, y):
+    name, dataset_fn = y
+    return x + combinations.combine(
+        dataset_fn=combinations.NamedObject(name, dataset_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class OptimizationTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptimizationStatefulFunction(self):
+    dataset = dataset_ops.Dataset.range(
+        10).map(lambda _: random_ops.random_uniform([])).batch(10)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    get_next = self.getNext(dataset)
+    self.evaluate(get_next())
+
+  # TODO(b/123354468)
+  @combinations.generate(test_base.graph_only_combinations())
+  def testOptimizationLargeInputFromTensor(self):
+    input_t = array_ops.placeholder(dtypes.int32, (None, None, None))
+    dataset = dataset_ops.Dataset.from_tensors(input_t)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    init_op = iterator.initializer
+    get_next = iterator.get_next()
+
+    with self.cached_session() as sess:
+      sess.run(init_op, {input_t: np.ones([512, 1024, 1025], np.int32)})
+      self.evaluate(get_next)
+
+  # TODO(b/123354468)
+  @combinations.generate(test_base.graph_only_combinations())
+  def testOptimizationLargeInputFromTensorSlices(self):
+    input_t = array_ops.placeholder(dtypes.int32, (None, None, None, None))
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_t)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    dataset = dataset.with_options(options)
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    init_op = iterator.initializer
+    get_next = iterator.get_next()
+
+    with self.cached_session() as sess:
+      sess.run(init_op, {input_t: np.ones([1, 512, 1024, 1025], np.int32)})
+      self.evaluate(get_next)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptimizationNestedDataset(self):
+
+    def flat_map_fn(_):
+      dataset = dataset_ops.Dataset.from_tensors(0)
+      dataset = dataset.apply(testing.assert_next(["MemoryCacheImpl"]))
+      dataset = dataset.skip(0)  # Should be removed by noop elimination
+      dataset = dataset.cache()
+      return dataset
+
+    dataset = dataset_ops.Dataset.range(1)
+    dataset = dataset.flat_map(flat_map_fn)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.noop_elimination = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[0])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOptimizationNestedDatasetWithModifiedRetval(self):
+
+    def flat_map_fn(_):
+      dataset = dataset_ops.Dataset.from_tensors(0)
+      dataset = dataset.apply(testing.assert_next(["MapAndBatch"]))
+      # Should be fused by map and batch fusion
+      dataset = dataset.map(lambda x: x)
+      dataset = dataset.batch(1)
+      return dataset
+
+    dataset = dataset_ops.Dataset.range(1)
+    dataset = dataset.flat_map(flat_map_fn)
+
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.map_and_batch_fusion = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, expected_output=[[0]])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(autotune=[True, False, None]),
+          combinations.combine(map_parallelization=[True, False, None])))
+  def testOptimizationMapParallelization(self, autotune, map_parallelization):
+    dataset = dataset_ops.Dataset.range(5)
+    if autotune is not False and map_parallelization is not False:  # pylint: disable=g-bool-id-comparison
+      dataset = dataset.apply(testing.assert_next(["ParallelMap"]))
+    else:
+      dataset = dataset.apply(testing.assert_next(["Map"]))
+    dataset = dataset.map(lambda x: x + 1)
+
+    options = options_lib.Options()
+    if autotune is not None:
+      options.autotune.enabled = autotune
+    if map_parallelization is not None:
+      options.experimental_optimization.map_parallelization = (
+          map_parallelization)
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(dataset, expected_output=list(range(1, 6)))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(existing_prefetch=[True, False]),
+                         combinations.combine(autotune=[True, False]),
+                         combinations.combine(inject_prefetch=[True, False])))
+  def testOptimizationInjectPrefetch(self, existing_prefetch, autotune,
+                                     inject_prefetch):
+    dataset = dataset_ops.Dataset.range(5)
+    dataset = dataset.map(
+        lambda x: x + 1, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.batch(1)
+    if existing_prefetch:
+      dataset = dataset.prefetch(1)
+    if autotune and inject_prefetch and not existing_prefetch:
+      dataset = dataset.apply(testing.assert_next(["Prefetch", "Root"]))
+    else:
+      dataset = dataset.apply(testing.assert_next(["Root"]))
+
+    options = options_lib.Options()
+    options.autotune.enabled = autotune
+    options.experimental_optimization.map_and_batch_fusion = False
+    if not inject_prefetch:
+      options.experimental_optimization.inject_prefetch = False
+    dataset = dataset.with_options(options)
+
+    self.assertDatasetProduces(dataset, expected_output=[np.array([x]) for x in
+                                                         range(1, 6)])
+
+  # Reference variables are not supported in eager mode.
+  @combinations.generate(
+      combinations.times(test_base.graph_only_combinations(),
+                         _captured_refvar_test_combinations()))
+  def testOptimizationWithCapturedRefVar(self, dataset_fn):
+    """Tests that default optimizations are disabled with ref variables."""
+    variable = variable_scope.get_variable(
+        "v", initializer=0, use_resource=False)
+    assign_op = variable.assign_add(1)
+    unoptimized_dataset = dataset_fn(variable)
+
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.noop_elimination = True
+    options.experimental_optimization.map_and_batch_fusion = True
+    options.experimental_warm_start = False
+    optimized_dataset = unoptimized_dataset.with_options(options)
+    optimized_it = dataset_ops.make_initializable_iterator(optimized_dataset)
+
+    # Check that outputs are the same in the optimized and unoptimized cases,
+    # when the variable value is changing.
+    unoptimized_it = dataset_ops.make_initializable_iterator(
+        unoptimized_dataset)
+    with ops.control_dependencies([assign_op]):
+      unoptimized_output = unoptimized_it.get_next()
+      optimized_output = optimized_it.get_next()
+
+    self.evaluate(variable.initializer)
+    self.evaluate((unoptimized_it.initializer, optimized_it.initializer))
+    while True:
+      try:
+        unoptimized, optimized = self.evaluate((unoptimized_output,
+                                                optimized_output))
+        self.assertEqual(unoptimized, optimized)
+      except errors.OutOfRangeError:
+        break
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(warm_start=[True, False]),
+      )
+  )
+  def testOptimizationWarmStart(self, warm_start):
+    dataset = dataset_ops.Dataset.range(10)
+    counter = variables.Variable(0)
+
+    def update_counter(x):
+      counter.assign_add(1)
+      return x
+
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    if warm_start:
+      options.experimental_warm_start = True
+    else:
+      options.experimental_warm_start = False
+    dataset = dataset.with_options(options)
+    dataset = dataset.map(update_counter).prefetch(10)
+    unused_iter = iter(dataset)
+
+    if warm_start:
+      for sleep_time_secs in [0.1, 0.2, 0.5, 2, 5, 10]:
+        if counter.numpy() == 0:
+          time.sleep(sleep_time_secs)
+        else:
+          break
+      self.assertGreater(counter.numpy(), 0)
+    else:
+      self.assertEqual(counter.numpy(), 0)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/seq_interleave_prefetch_test.py

@@ -0,0 +1,100 @@
+# Copyright 2023 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 `SeqInterleavePrefetch` optimization."""
+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 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.platform import test
+
+
+class SeqInterleavePrefetchTest(
+    test_base.DatasetTestBase, parameterized.TestCase
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(cycle_length=[2, 4]),
+          combinations.combine(block_length=[2, 4]),
+          combinations.combine(other_arguments=[True, False]),
+      )
+  )
+  def testOptimizationSeqInterleavePrefetch(
+      self,
+      cycle_length,
+      block_length,
+      other_arguments,
+  ):
+    num_input_elements = 16
+    var1 = constant_op.constant(9, dtype=dtypes.int64)
+    var2 = constant_op.constant(11, dtype=dtypes.int64)
+
+    # dataset1: Deterministic parallel interleave dataset.
+    dataset1 = dataset_ops.Dataset.range(num_input_elements)
+    options1 = options_lib.Options()
+    options1.experimental_optimization.apply_default_optimizations = False
+    options1.experimental_optimization.seq_interleave_prefetch = False
+    dataset1 = dataset1.with_options(options1)
+    if other_arguments:
+      dataset1 = dataset1.interleave(
+          (lambda _: dataset_ops.Dataset.range(var1 + var2 + 1)),
+          cycle_length=cycle_length,
+          block_length=block_length,
+          num_parallel_calls=dataset_ops.AUTOTUNE,
+          deterministic=True,
+      )
+    else:
+      dataset1 = dataset1.interleave(
+          (lambda _: dataset_ops.Dataset.range(num_input_elements)),
+          cycle_length=cycle_length,
+          block_length=block_length,
+          num_parallel_calls=dataset_ops.AUTOTUNE,
+          deterministic=True,
+      )
+
+    # dataset2: Deterministic parallel interleave dataset with
+    # `seq_interleave_prefetch` optimization enabled.
+    dataset2 = dataset_ops.Dataset.range(num_input_elements)
+    options2 = options_lib.Options()
+    options2.experimental_optimization.apply_default_optimizations = False
+    options2.experimental_optimization.seq_interleave_prefetch = True
+    dataset2 = dataset2.with_options(options2)
+    if other_arguments:
+      dataset2 = dataset2.interleave(
+          (lambda _: dataset_ops.Dataset.range(var1 + var2 + 1)),
+          cycle_length=cycle_length,
+          block_length=block_length,
+          num_parallel_calls=dataset_ops.AUTOTUNE,
+          deterministic=True,
+      )
+    else:
+      dataset2 = dataset2.interleave(
+          (lambda _: dataset_ops.Dataset.range(num_input_elements)),
+          cycle_length=cycle_length,
+          block_length=block_length,
+          num_parallel_calls=dataset_ops.AUTOTUNE,
+          deterministic=True,
+      )
+
+    self.assertDatasetsEqual(dataset1, dataset2)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/optimization/shuffle_and_repeat_fusion_test.py

@@ -0,0 +1,53 @@
+# 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 `ShuffleAndRepeatFusion` optimization."""
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class ShuffleAndRepeatFusionTest(test_base.DatasetTestBase,
+                                 parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleAndRepeatFusion(self):
+    expected = "ShuffleAndRepeat"
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.assert_next([expected])).shuffle(10).repeat(2)
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    options.experimental_optimization.shuffle_and_repeat_fusion = True
+    dataset = dataset.with_options(options)
+    get_next = self.getNext(dataset)
+
+    for _ in range(2):
+      results = []
+      for _ in range(10):
+        results.append(self.evaluate(get_next()))
+      self.assertAllEqual([x for x in range(10)], sorted(results))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,124 @@
+# Copyright 2022 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.pad_to_cardinality()."""
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import pad_to_cardinality
+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
+
+
+pad_to_cardinality = pad_to_cardinality.pad_to_cardinality
+
+
+class PadToCardinalityTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    data = [1, 2, 3, 4, 5]
+    target = 12
+    ds = dataset_ops.Dataset.from_tensor_slices({'a': data})
+    ds = ds.apply(pad_to_cardinality(target))
+    expected_data = [{'a': data[i], 'valid': True} for i in range(len(data))]
+    expected_padding = [
+        {'a': 0, 'valid': False} for _ in range(target - len(data))
+    ]
+    self.assertAllEqual(
+        self.getDatasetOutput(ds), expected_data + expected_padding
+    )
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoPadding(self):
+    data = [1, 2, 3, 4, 5]
+    target = 5
+    ds = dataset_ops.Dataset.from_tensor_slices({'a': data})
+    ds = ds.apply(pad_to_cardinality(target))
+    expected_data = [{'a': data[i], 'valid': True} for i in range(len(data))]
+    self.assertAllEqual(self.getDatasetOutput(ds), expected_data)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStructuredData(self):
+    data = {
+        'a': [1, 2, 3, 4, 5],
+        'b': ([b'a', b'b', b'c', b'd', b'e'], [-1, -2, -3, -4, -5]),
+    }
+    data_len = len(data['a'])
+    target = 12
+    ds = dataset_ops.Dataset.from_tensor_slices(data)
+    ds = ds.apply(pad_to_cardinality(target))
+    expected_data = [
+        {
+            'a': data['a'][i],
+            'b': (data['b'][0][i], data['b'][1][i]),
+            'valid': True,
+        }
+        for i in range(data_len)
+    ]
+    expected_padding = [
+        {'a': 0, 'b': (b'', 0), 'valid': False}
+        for _ in range(target - data_len)
+    ]
+    self.assertAllEqual(
+        self.getDatasetOutput(ds), expected_data + expected_padding
+    )
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testUnknownCardinality(self):
+    ds = dataset_ops.Dataset.from_tensors({'a': 1}).filter(lambda x: True)
+    with self.assertRaisesRegex(
+        ValueError,
+        'The dataset passed into `pad_to_cardinality` must have a '
+        'known cardinalty, but has cardinality -2',
+    ):
+      ds = ds.apply(pad_to_cardinality(5))
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testInfiniteCardinality(self):
+    ds = dataset_ops.Dataset.from_tensors({'a': 1}).repeat()
+    with self.assertRaisesRegex(
+        ValueError,
+        'The dataset passed into `pad_to_cardinality` must have a '
+        'known cardinalty, but has cardinality -1',
+    ):
+      ds = ds.apply(pad_to_cardinality(5))
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testNonMapping(self):
+    ds = dataset_ops.Dataset.from_tensors(1)
+    with self.assertRaisesRegex(
+        ValueError,
+        '`pad_to_cardinality` requires its input dataset to be a dictionary',
+    ):
+      ds = ds.apply(pad_to_cardinality(5))
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.v2_eager_only_combinations())
+  def testRequestedCardinalityTooShort(self):
+    ds = dataset_ops.Dataset.from_tensors({'a': 1}).repeat(5)
+    with self.assertRaisesRegex(
+        ValueError,
+        r'The dataset passed into `pad_to_cardinality` must have a cardinalty '
+        r'less than the target cardinality \(3\), but has cardinality 5',
+    ):
+      ds = ds.apply(pad_to_cardinality(3))
+      self.getDatasetOutput(ds)
+
+
+if __name__ == '__main__':
+  test.main()

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

@@ -0,0 +1,460 @@
+# 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.parallel_interleave()`."""
+import itertools
+import math
+import threading
+import time
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import interleave_ops
+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 dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.platform import test
+
+
+class ParallelInterleaveTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+
+    self.error = None
+    self.repeat_count = 2
+
+    # Set up threading events used to sequence when items are produced that
+    # are subsequently interleaved. These events allow us to deterministically
+    # simulate slowdowns and force sloppiness.
+    self.read_coordination_events = {}
+    self.write_coordination_events = {}
+    # input values [4, 5, 6] are the common case for the tests; set defaults
+    for i in range(4, 7):
+      self.read_coordination_events[i] = threading.Semaphore(0)
+      self.write_coordination_events[i] = threading.Event()
+
+  def dataset_fn(self, input_values, cycle_length, block_length, sloppy,
+                 buffer_output_elements, prefetch_input_elements):
+
+    def map_py_fn(x):
+      self.write_coordination_events[x].wait()
+      self.write_coordination_events[x].clear()
+      self.read_coordination_events[x].release()
+      if self.error:
+        err = self.error
+        self.error = None
+        raise err  # pylint: disable=raising-bad-type
+      return x * x
+
+    def map_fn(x):
+      return script_ops.py_func(map_py_fn, [x], x.dtype)
+
+    def interleave_fn(x):
+      dataset = dataset_ops.Dataset.from_tensors(x)
+      dataset = dataset.repeat(x)
+      return dataset.map(map_fn)
+
+    return dataset_ops.Dataset.from_tensor_slices(input_values).repeat(
+        self.repeat_count).apply(
+            interleave_ops.parallel_interleave(
+                interleave_fn, cycle_length, block_length, sloppy,
+                buffer_output_elements, prefetch_input_elements))
+
+  def _interleave(self, lists, cycle_length, block_length):
+    """Python implementation of interleave used for testing."""
+    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 = []
+    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
+
+  @combinations.generate(
+      combinations.times(
+          combinations.combine(
+              input_lists=[[[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]]],
+              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
+              ]],
+              cycle_length=1,
+              block_length=1) +
+          combinations.combine(
+              input_lists=[[[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]]],
+              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
+              ]],
+              cycle_length=2,
+              block_length=1) + combinations.combine(
+                  input_lists=[[[4] * 4, [5] * 5, [6] * 6] * 2],
+                  expected_elements=[[
+                      4, 4, 5, 5, 4, 4, 5, 5, 5, 6, 6, 4, 4, 6, 6, 4, 4, 6, 6,
+                      5, 5, 6, 6, 5, 5, 6, 6, 5, 6, 6
+                  ]],
+                  cycle_length=2,
+                  block_length=2) +
+          combinations.combine(
+              input_lists=[[[4, 4, 4, 4], [], [6, 6, 6, 6, 6, 6], [4, 4, 4, 4],
+                            [], [6, 6, 6, 6, 6, 6]]],
+              expected_elements=[[
+                  4, 4, 6, 4, 6, 4, 6, 6, 4, 6, 4, 6, 4, 4, 6, 6, 6, 6, 6, 6
+              ]],
+              cycle_length=2,
+              block_length=1)))
+  def testPythonImplementation(self, input_lists, expected_elements,
+                               cycle_length, block_length):
+    for index, (expected, produced) in enumerate(
+        itertools.zip_longest(
+            expected_elements,
+            self._interleave(input_lists, cycle_length, block_length))):
+      self.assertEqual(expected, produced, "Values differ at %s. %s != %s" %
+                       (index, expected, produced))
+
+  def _clear_coordination_events(self):
+    for i in range(4, 7):
+      self.read_coordination_events[i] = threading.Semaphore(0)
+      self.write_coordination_events[i].clear()
+
+  def _allow_all_map_threads(self):
+    for i in range(4, 7):
+      self.write_coordination_events[i].set()
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(sloppy=[False, True])))
+  def testEmptyInput(self, sloppy):
+    # Empty input.
+    self._clear_coordination_events()
+    next_element = self.getNext(
+        self.dataset_fn(
+            input_values=np.int64([]),
+            cycle_length=2,
+            block_length=3,
+            sloppy=sloppy,
+            buffer_output_elements=1,
+            prefetch_input_elements=0))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(sloppy=[False, True])))
+  def _testNonEmptyInputIntoEmptyOutputs(self, sloppy):
+    # Non-empty input leading to empty output.
+    self._clear_coordination_events()
+    next_element = self.getNext(
+        self.dataset_fn(
+            input_values=np.int64([0, 0, 0]),
+            cycle_length=2,
+            block_length=3,
+            sloppy=sloppy,
+            buffer_output_elements=1,
+            prefetch_input_elements=0))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(sloppy=[False, True])))
+  def testTooManyReaders(self, sloppy=False):
+
+    def interleave_fn(x):
+      dataset = dataset_ops.Dataset.from_tensors(x)
+      dataset = dataset.repeat(math_ops.cast(x, dtype=dtypes.int64))
+      return dataset
+
+    dataset = dataset_ops.Dataset.from_tensor_slices([4, 5, 6])
+    dataset = dataset.repeat(self.repeat_count)
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(
+            interleave_fn, cycle_length=16, block_length=2, sloppy=sloppy))
+    get_next = self.getNext(dataset)
+    output_values = []
+    for _ in range(30):
+      output_values.append(self.evaluate(get_next()))
+
+    expected_values = self._interleave(
+        [[4] * 4, [5] * 5, [6] * 6] * self.repeat_count, 1, 2)
+    self.assertCountEqual(output_values, expected_values)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSparse(self):
+    def _map_fn(i):
+      return sparse_tensor.SparseTensor(
+          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).apply(
+        interleave_ops.parallel_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())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorsInInputFn(self):
+
+    def map_py_fn(x):
+      if x == 5:
+        raise ValueError()
+      return x
+
+    def map_fn(x):
+      return script_ops.py_func(map_py_fn, [x], x.dtype)
+
+    def interleave_fn(x):
+      dataset = dataset_ops.Dataset.from_tensors(x)
+      dataset = dataset.repeat(x)
+      return dataset
+
+    def dataset_fn(input_values, cycle_length, block_length, sloppy,
+                   buffer_output_elements, prefetch_input_elements):
+      return dataset_ops.Dataset.from_tensor_slices(input_values).map(
+          map_fn).repeat(self.repeat_count).apply(
+              interleave_ops.parallel_interleave(
+                  interleave_fn, cycle_length, block_length, sloppy,
+                  buffer_output_elements, prefetch_input_elements))
+
+    next_element = self.getNext(
+        dataset_fn(
+            input_values=np.int64([4, 5, 6]),
+            cycle_length=2,
+            block_length=1,
+            sloppy=False,
+            buffer_output_elements=1,
+            prefetch_input_elements=0))
+    for i, expected_element in enumerate(
+        self._interleave([[4] * 4, [5], [6] * 6] * self.repeat_count, 2, 1)):
+      if expected_element == 5:
+        with self.assertRaises(errors.InvalidArgumentError):
+          self.evaluate(next_element())
+      else:
+        actual_element = self.evaluate(next_element())
+        self.assertEqual(
+            expected_element, actual_element,
+            "At index %s: %s expected, got: %s" % (i, expected_element,
+                                                   actual_element))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testErrorsInInterleaveFn(self):
+
+    def map_py_fn(x):
+      if x == 5:
+        raise ValueError()
+      return x
+
+    def interleave_fn(x):
+      dataset = dataset_ops.Dataset.from_tensors(x)
+      y = script_ops.py_func(map_py_fn, [x], x.dtype)
+      dataset = dataset.repeat(y)
+      return dataset
+
+    def dataset_fn(input_values, cycle_length, block_length, sloppy,
+                   buffer_output_elements, prefetch_input_elements):
+      return dataset_ops.Dataset.from_tensor_slices(input_values).repeat(
+          self.repeat_count).apply(
+              interleave_ops.parallel_interleave(
+                  interleave_fn, cycle_length, block_length, sloppy,
+                  buffer_output_elements, prefetch_input_elements))
+
+    next_element = self.getNext(
+        dataset_fn(
+            input_values=np.int64([4, 5, 6]),
+            cycle_length=2,
+            block_length=1,
+            sloppy=False,
+            buffer_output_elements=1,
+            prefetch_input_elements=0))
+    for i, expected_element in enumerate(
+        self._interleave([[4] * 4, [5], [6] * 6] * self.repeat_count, 2, 1)):
+      if expected_element == 5:
+        with self.assertRaises(errors.InvalidArgumentError):
+          self.evaluate(next_element())
+      else:
+        actual_element = self.evaluate(next_element())
+        self.assertEqual(
+            expected_element, actual_element,
+            "At index %s: %s expected, got: %s" % (i, expected_element,
+                                                   actual_element))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShutdownRace(self):
+    dataset = dataset_ops.Dataset.range(20)
+    map_fn = lambda x: dataset_ops.Dataset.range(20 * x, 20 * (x + 1))
+    dataset = dataset.apply(
+        interleave_ops.parallel_interleave(
+            map_fn,
+            cycle_length=3,
+            sloppy=False,
+            buffer_output_elements=1,
+            prefetch_input_elements=0))
+    dataset = dataset.batch(32)
+
+    results = []
+    for _ in range(2):
+      elements = []
+      next_element = self.getNext(dataset)
+      try:
+        while True:
+          elements.extend(self.evaluate(next_element()))
+      except errors.OutOfRangeError:
+        pass
+      results.append(elements)
+    self.assertAllEqual(results[0], results[1])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              sloppy=[None, True, False], global_determinism=[True, False])))
+  def testDeterminismConfiguration(self, sloppy, global_determinism):
+    if sloppy is None:
+      expect_determinism = global_determinism
+    else:
+      expect_determinism = not sloppy
+    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.apply(
+          interleave_ops.parallel_interleave(
+              interleave_fn, cycle_length=10, sloppy=sloppy))
+
+      opts = options_lib.Options()
+      opts.deterministic = global_determinism
+      dataset = dataset.with_options(opts)
+      return dataset
+
+    self.checkDeterminism(dataset_fn, expect_determinism, elements)
+
+
+class ParallelInterleaveCheckpointTest(
+    checkpoint_test_base.CheckpointTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(ParallelInterleaveCheckpointTest, self).setUp()
+    self.input_values = np.array([2, 3], dtype=np.int64)
+    self.num_repeats = 2
+    self.num_outputs = np.sum(self.input_values) * 2
+
+  def _build_ds(self, cycle_length, block_length, sloppy=False):
+    return (dataset_ops.Dataset.from_tensor_slices(self.input_values).repeat(
+        self.num_repeats).apply(
+            interleave_ops.parallel_interleave(
+                lambda x: dataset_ops.Dataset.range(10 * x, 11 * x),
+                cycle_length, block_length, sloppy)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(cycle_length=[1, 2], block_length=[1, 3])))
+  def test(self, verify_fn, cycle_length, block_length):
+    verify_fn(self, lambda: self._build_ds(cycle_length, block_length),
+              self.num_outputs)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(cycle_length=[1, 2], block_length=[1, 3])))
+  def testWithSloppy(self, cycle_length, block_length):
+    break_points = self.gen_break_points(self.num_outputs, 10)
+    expected_outputs = np.repeat(
+        np.concatenate([np.arange(10 * x, 11 * x) for x in self.input_values]),
+        self.num_repeats).tolist()
+
+    actual = self.gen_outputs(
+        lambda: self._build_ds(cycle_length, block_length, True),
+        break_points, self.num_outputs)
+    self.assertSequenceEqual(sorted(actual), expected_outputs)
+
+  @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).apply(
+          interleave_ops.parallel_interleave(_interleave_fn, 1))
+
+    verify_fn(self, _build_dataset, num_outputs=20)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,225 @@
+# 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.prefetch_to_device()`."""
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import config_pb2
+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.data.util import structure
+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 test_util
+from tensorflow.python.platform import test
+
+
+# TODO(b/117581999): add eager coverage when supported.
+class PrefetchToDeviceTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchToDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(structure.are_compatible(
+        dataset_ops.get_structure(host_dataset),
+        dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchToSameDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device(
+            "/job:localhost/replica:0/task:0/device:CPU:0"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(structure.are_compatible(
+        dataset_ops.get_structure(host_dataset),
+        dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchDictToDevice(self):
+    host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x})
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(structure.are_compatible(
+        dataset_ops.get_structure(host_dataset),
+        dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element["a"].dtype)
+    self.assertEqual([], next_element["a"].shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        self.assertEqual({"a": i}, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchSparseTensorsToDevice(self):
+    def make_tensor(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=[[0, 0]], values=(i*[1]), dense_shape=[2, 2])
+    host_dataset = dataset_ops.Dataset.range(10).map(make_tensor)
+
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_one_shot_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(structure.are_compatible(
+        dataset_ops.get_structure(host_dataset),
+        dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      for i in range(10):
+        actual = self.evaluate(next_element)
+        self.assertAllEqual([i], actual.values)
+        self.assertAllEqual([[0, 0]], actual.indices)
+        self.assertAllEqual([2, 2], actual.dense_shape)
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPrefetchToDeviceGpu(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/gpu:0"))
+
+    self.assertDatasetProduces(device_dataset, list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPrefetchToDeviceCorrectPlacement(self):
+
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
+
+    self.assertIn("gpu:0", dataset._variant_tensor.device.lower())
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchToDeviceWithReInit(self):
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/cpu:1"))
+
+    with ops.device("/cpu:1"):
+      iterator = dataset_ops.make_initializable_iterator(device_dataset)
+      next_element = iterator.get_next()
+
+    self.assertTrue(structure.are_compatible(
+        dataset_ops.get_structure(host_dataset),
+        dataset_ops.get_structure(device_dataset)))
+
+    self.assertEqual(dtypes.int64, next_element.dtype)
+    self.assertEqual([], next_element.shape)
+
+    worker_config = config_pb2.ConfigProto(device_count={"CPU": 2})
+    with self.test_session(config=worker_config):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPrefetchToDeviceGpuWithReInit(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/gpu:0"))
+
+    iterator = dataset_ops.make_initializable_iterator(device_dataset)
+    next_element = iterator.get_next()
+
+    with self.cached_session(
+        config=config_pb2.ConfigProto(allow_soft_placement=False)):
+      self.evaluate(iterator.initializer)
+      for i in range(5):
+        self.assertEqual(i, self.evaluate(next_element))
+      self.evaluate(iterator.initializer)
+      for i in range(10):
+        self.assertEqual(i, self.evaluate(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testPrefetchToDevicePlacement(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("No GPU available")
+
+    host_dataset = dataset_ops.Dataset.range(10)
+    device_dataset = host_dataset.apply(
+        prefetching_ops.prefetch_to_device("/gpu:0"))
+
+    self.assertEqual(device_dataset._variant_tensor.device,
+                     "/job:localhost/replica:0/task:0/device:GPU:0")
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,97 @@
+# 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 `experimental_slack` option."""
+
+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.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 PrefetchWithSlackTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(PrefetchWithSlackTest, self).setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(3)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPrefetchWithSlackOption(self):
+    """Determines slack_period based on num devices attached to iterator."""
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.prefetch(1)
+    options = options_lib.Options()
+    options.experimental_slack = 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)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPrefetchWithSlackOptionWithoutIterator(self):
+    """Defaults to slack period of 1 without iterator."""
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.prefetch(1)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, range(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWithPassthroughDataset(self):
+    """Should still work with a passthrough dataset after prefetch()."""
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.prefetch(1)
+    dataset = dataset.map(lambda x: x + 1)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, range(1, 11))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoErrorWithoutPrefetch(self):
+    """The rewrite should not fail if there is no prefetch() in the pipeline."""
+    dataset = dataset_ops.Dataset.range(10)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, range(10))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoErrorWithInvalidDataset(self):
+    """With a nested dataset op after prefetch, the rewrite should fail."""
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.prefetch(1)
+    dataset = dataset.flat_map(dataset_ops.Dataset.from_tensors)
+    options = options_lib.Options()
+    options.experimental_slack = True
+    dataset = dataset.with_options(options)
+    self.assertDatasetProduces(dataset, range(10))
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,412 @@
+# 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 private `_RebatchDataset` transformation."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import distribute
+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 constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+class BatchSizesForWorkerTest(test_base.DatasetTestBase,
+                              parameterized.TestCase):
+
+  def _test(self, global_batch_size, num_workers, num_replicas_per_worker,
+            is_batch_size_static):
+    """Test that all constraints are met for given parameters."""
+    if not is_batch_size_static:
+      # Adding a constant value here prevents downstream computation from
+      # statically deriving the value of global batch size when running
+      # in graph mode.
+      global_batch_size += constant_op.constant(0, dtypes.int64)
+
+    batch_sizes_list = []
+    for i in range(num_workers):
+      batch_sizes_list.append(
+          self.evaluate(
+              distribute.batch_sizes_for_worker(global_batch_size, num_workers,
+                                                num_replicas_per_worker, i)))
+    for batch_sizes in batch_sizes_list:
+      # Constraint (A): for any worker, len(batch_sizes) == W * R
+      self.assertLen(batch_sizes, num_workers * num_replicas_per_worker)
+      # Constraint (B): for any worker, sum(batch_sizes) == G
+      self.assertAllEqual(np.sum(batch_sizes), global_batch_size)
+
+    # Each per-worker batch is split into num_workers global steps
+    for step_index in range(num_workers):
+      actual_global_batch = 0
+      offset = step_index * num_replicas_per_worker
+      for batch_sizes in batch_sizes_list:
+        actual_global_batch += np.sum(batch_sizes[offset:offset +
+                                                  num_replicas_per_worker])
+      # Constraint (C): for any step, batch size across all workers add up to G.
+      self.assertAllEqual(
+          global_batch_size,
+          actual_global_batch,
+      )
+
+    # Constraint (D): Batch size of any two replicas differs by at most one
+    self.assertLessEqual(np.max(batch_sizes_list) - np.min(batch_sizes_list), 1)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(is_batch_size_static=[True, False])))
+  def testBasic(self, is_batch_size_static):
+    # Manually verify basic test case.
+    global_batch_size = 8
+    num_workers = 2
+    num_replicas_per_worker = 2
+    for worker_index in range(4):
+      batch_sizes = distribute.batch_sizes_for_worker(global_batch_size,
+                                                      num_workers,
+                                                      num_replicas_per_worker,
+                                                      worker_index)
+      self.assertAllEqual([2, 2, 2, 2],
+                          tensor_util.constant_value(batch_sizes))
+    self._test(global_batch_size, num_workers, num_replicas_per_worker,
+               is_batch_size_static)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(is_batch_size_static=[True, False])))
+  def testBatchSizeIndivisibleByNumWorkers(self, is_batch_size_static):
+    global_batch_size = 4
+    num_workers = 3
+    num_replicas_per_worker = 1
+
+    def get_batch_sizes_for_worker(worker_index):
+      return tensor_util.constant_value(
+          distribute.batch_sizes_for_worker(global_batch_size, num_workers,
+                                            num_replicas_per_worker,
+                                            worker_index))
+
+    # Manually verify this test case.
+    self.assertAllEqual([2, 1, 1], get_batch_sizes_for_worker(0))
+    self.assertAllEqual([1, 1, 2], get_batch_sizes_for_worker(1))
+    self.assertAllEqual([1, 2, 1], get_batch_sizes_for_worker(2))
+    self._test(global_batch_size, num_workers, num_replicas_per_worker,
+               is_batch_size_static)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(is_batch_size_static=[True, False])))
+  def testBatchSizeIndivisibleByNumReplicas(self, is_batch_size_static):
+    self._test(
+        global_batch_size=4,
+        num_workers=1,
+        num_replicas_per_worker=5,
+        is_batch_size_static=is_batch_size_static)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(is_batch_size_static=[True, False])))
+  def testBatchSizeSmallerThanNumReplicas(self, is_batch_size_static):
+    self._test(
+        global_batch_size=4,
+        num_workers=2,
+        num_replicas_per_worker=5,
+        is_batch_size_static=is_batch_size_static)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(is_batch_size_static=[True, False])))
+  def testBatchSizeSmallerThanNumWorkers(self, is_batch_size_static):
+    self._test(
+        global_batch_size=4,
+        num_workers=5,
+        num_replicas_per_worker=1,
+        is_batch_size_static=is_batch_size_static)
+
+
+def _flat_shapes(dataset):
+  return [
+      ts.as_list()
+      for ts in nest.flatten(dataset_ops.get_legacy_output_shapes(dataset))
+  ]
+
+
+class LegacyRebatchDatasetTest(test_base.DatasetTestBase,
+                               parameterized.TestCase):
+
+  @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=drop_remainder)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=2)
+
+    expected_shapes = [[2]] if drop_remainder else [[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(test_base.default_test_combinations())
+  def testCanHandleUnknownRank(self):
+    dataset = dataset_ops.Dataset.from_tensors("xxx")
+    # decode_image results in a tensor of completely unknown shape (i.e. unknown
+    # rank)
+    dataset = dataset.map(image_ops.decode_image)
+    self.assertEqual([tensor_shape.TensorShape(None)],
+                     nest.flatten(
+                         dataset_ops.get_legacy_output_shapes(dataset)))
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=4)
+    # Note that we are just testing the dataset shapes, not the actual output.
+    self.assertEqual(
+        [tensor_shape.TensorShape(None)],
+        nest.flatten(dataset_ops.get_legacy_output_shapes(rebatched_dataset)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCanHandleUnknownDims(self):
+    dataset = dataset_ops.Dataset.range(1000)
+    dataset = dataset.batch(10, drop_remainder=False)
+    dataset = dataset.batch(10, drop_remainder=False)
+    self.assertEqual([[None, None]], _flat_shapes(dataset))
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=4)
+    # Note that we are just testing the dataset shapes, not the actual output.
+    self.assertEqual([[None, None]], _flat_shapes(rebatched_dataset))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testScalarInputError(self):
+    dataset = dataset_ops.Dataset.range(1024)
+    distribute._LegacyRebatchDataset(dataset.batch(4), num_replicas=4)
+    with self.assertRaises(ValueError):
+      distribute._LegacyRebatchDataset(dataset, num_replicas=4)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testBatchNotDivisibleByNumReplicas(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(8).batch(
+        4, drop_remainder=drop_remainder)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=3)
+    self.assertEqual([[None]], _flat_shapes(rebatched_dataset))
+    # This rebatches into sub-batches of size 2, since ceil(4 / 3) = 2. However,
+    # this means that only the first 2 replicas will get data.
+    expected_output = [[0, 1], [2, 3], [], [4, 5], [6, 7], []]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTupleOutput(self):
+    dataset = dataset_ops.Dataset.range(1024).map(lambda x: (x, x)).batch(32)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=4)
+    expected_output = [([k for k in range(i, i + 8)],  # pylint: disable=g-complex-comprehension
+                        [k for k in range(i, i + 8)])
+                       for i in range(0, 1024, 8)]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedDictionaryOutput(self):
+    dataset = dataset_ops.Dataset.range(8).map(
+        lambda x: {"a": x, "b": {"c": x + 1}}).batch(4)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=2)
+    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 testFinalPartialBatch(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(10).batch(
+        4, drop_remainder=drop_remainder)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=2)
+    self.assertEqual([[2] if drop_remainder else [None]],
+                     _flat_shapes(rebatched_dataset))
+    if drop_remainder:
+      expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    else:
+      expected_output = [[0, 1], [2, 3], [4, 5], [6, 7], [8], [9]]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(drop_remainder=[True, False])))
+  def testFinalPartialBatchAfterRebatch(self, drop_remainder):
+    dataset = dataset_ops.Dataset.range(9).batch(
+        4, drop_remainder=drop_remainder)
+    rebatched_dataset = distribute._LegacyRebatchDataset(
+        dataset, num_replicas=2)
+    self.assertEqual([[2] if drop_remainder else [None]],
+                     _flat_shapes(rebatched_dataset))
+    if drop_remainder:
+      expected_output = [[0, 1], [2, 3], [4, 5], [6, 7]]
+    else:
+      expected_output = [[0, 1], [2, 3], [4, 5], [6, 7], [8], []]
+    self.assertDatasetProduces(rebatched_dataset, expected_output)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleBatches(self):
+    dataset = dataset_ops.Dataset.range(16).batch(2).batch(4)
+    self.assertEqual([[None, None]], _flat_shapes(dataset))
+
+    # Each element is a list of 4 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(dataset, expected_output)
+
+    rebatched_dataset = distribute._LegacyRebatchDataset(dataset, 2)
+    self.assertEqual([[None, None]], _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 testRaggedTensorDataset(self):
+    # Set up a dataset that produces ragged tensors with a static batch size.
+    row_lengths = np.random.randint(8, size=128)
+    values = np.random.normal(size=np.sum(row_lengths)).astype(np.float32)
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        ragged_tensor.RaggedTensor.from_row_lengths(values, row_lengths))
+    dataset = dataset.batch(32, drop_remainder=True)
+
+    # The map changes the internal representation of the ragged tensor.
+    # This test will fail if we don't normalize the tensor representation.
+    dataset = dataset.map(lambda x: x)
+
+    dataset = distribute._LegacyRebatchDataset(dataset, num_replicas=8)
+    # After rebatching, batch size is now 4.
+    expected_output = []
+    value_index = 0
+    for batch_row_lengths in row_lengths.reshape((-1, 4)):
+      num_values = np.sum(batch_row_lengths)
+      expected_output.append(
+          ragged_tensor.RaggedTensor.from_row_lengths(
+              values[value_index:(value_index + num_values)],
+              batch_row_lengths))
+      value_index += num_values
+    self.assertDatasetProduces(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)
+    _ = distribute._LegacyRebatchDataset(dataset, num_replicas=2)
+
+
+class ComputeBatchSizeTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeKnown(self):
+    # When drop_remainder=True, batch size can be inferred from the type spec.
+    dataset = dataset_ops.Dataset.range(32).batch(4, drop_remainder=True)
+    dataset = dataset_ops.Dataset.zip((dataset, dataset))
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(4, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeKnownAndMismatched(self):
+    # Return -1 when different components have different batch sizes.
+    dataset = dataset_ops.Dataset.range(32)
+    dataset = dataset_ops.Dataset.zip((dataset.batch(4, drop_remainder=True),
+                                       dataset.batch(8, drop_remainder=True)))
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(-1, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeUnknown(self):
+    dataset = dataset_ops.Dataset.range(32).batch(4)
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(4, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeWithPassthrough(self):
+    dataset = dataset_ops.Dataset.range(32).batch(4)
+    dataset = dataset.take(5)
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(4, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeWithPassthroughInvalid(self):
+    dataset = dataset_ops.Dataset.range(32).batch(4)
+    dataset = dataset.map(lambda x: x + 1)
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(-1, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeWithZip(self):
+    dataset = dataset_ops.Dataset.range(32).batch(4)
+    dataset = dataset_ops.Dataset.zip((dataset, dataset))
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(4, self.evaluate(batch_size))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testComputeBatchSizeWithZipMismatched(self):
+    dataset = dataset_ops.Dataset.range(32)
+    dataset = dataset_ops.Dataset.zip((dataset.batch(4), dataset.batch(8)))
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(-1, self.evaluate(batch_size))
+
+  @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 computing batch size logic.
+    dataset = dataset_ops.Dataset.range(4)
+    dataset = dataset.map(lambda x: (x, None))
+    dataset = dataset.batch(4, drop_remainder=True)
+    batch_size = distribute.compute_batch_size(dataset)
+    self.assertEqual(4, self.evaluate(batch_size))
+
+
+class LegacyRebatchDatasetCheckpointTest(
+    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):
+      return distribute._LegacyRebatchDataset(
+          dataset_ops.Dataset.range(num_elements).batch(
+              4 * batch_size, drop_remainder=True),
+          num_replicas=4)
+
+    verify_fn(self, lambda: build_dataset(64, 8), num_outputs=8)
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,428 @@
+# 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 private `replicate()` transformation."""
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import distribute
+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 config
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class LocalReplicateTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
+    super(LocalReplicateTest, self).__init__(methodName)
+    cpus = config.list_physical_devices("CPU")
+    # Set 3 virtual CPUs
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    self._device0 = "/device:CPU:0"
+    self._device1 = "/device:CPU:1"
+    self._device2 = "/device:CPU:2"
+
+  def tearDown(self):
+    super().tearDown()
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+
+    with ops.device(self._device0):
+      self.assertDatasetProduces(dataset0, range(100))
+    with ops.device(self._device1):
+      self.assertDatasetProduces(dataset1, range(100))
+    with ops.device(self._device2):
+      self.assertDatasetProduces(dataset2, range(100))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorsWithDataset(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100)
+      dataset0 = dataset_ops.Dataset.from_tensors(dataset0)
+      dataset0 = dataset0.flat_map(lambda x: x)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+
+    with ops.device(self._device0):
+      self.assertDatasetProduces(dataset0, range(100))
+    with ops.device(self._device1):
+      self.assertDatasetProduces(dataset1, range(100))
+    with ops.device(self._device2):
+      self.assertDatasetProduces(dataset2, range(100))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromTensorSlicesWithDataset(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100)
+      dataset0 = dataset_ops.Dataset.from_tensor_slices([dataset0])
+      dataset0 = dataset0.flat_map(lambda x: x)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+
+    with ops.device(self._device0):
+      self.assertDatasetProduces(dataset0, range(100))
+    with ops.device(self._device1):
+      self.assertDatasetProduces(dataset1, range(100))
+    with ops.device(self._device2):
+      self.assertDatasetProduces(dataset2, range(100))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testVariableInput(self):
+    with ops.device(self._device0):
+      counter_var = variable_scope.get_variable(
+          "counter", (), dtypes.int32, use_resource=True)
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: counter_var.assign_add(1))
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    self.evaluate(counter_var.initializer)
+    with ops.device(self._device1):
+      self.assertDatasetProduces(
+          dataset1, range(1, 101), requires_initialization=True)
+    with ops.device(self._device2):
+      self.assertDatasetProduces(
+          dataset2, range(1, 101), requires_initialization=True)
+    # Iterate through the original device last so that replication happens
+    # before counter_var is modified. The order only matters in graph mode.
+    with ops.device(self._device0):
+      self.assertDatasetProduces(
+          dataset0, range(1, 101), requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExternalStatePolicyIgnore(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: random_ops.random_uniform(  # pylint:disable=g-long-lambda
+              [],
+              minval=1,
+              maxval=10,
+              dtype=dtypes.float32))
+      opt = options_lib.Options()
+      opt.experimental_external_state_policy = (
+          options_lib.ExternalStatePolicy.IGNORE)
+      dataset0 = dataset0.with_options(opt)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+
+    with ops.device(self._device0):
+      get_next0 = self.getNext(dataset0)
+    with ops.device(self._device1):
+      get_next1 = self.getNext(dataset1)
+    with ops.device(self._device2):
+      get_next2 = self.getNext(dataset2)
+
+    for _ in range(100):
+      self.evaluate(get_next0())
+      self.evaluate(get_next1())
+      self.evaluate(get_next2())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExternalStatePolicyWarn(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: random_ops.random_uniform(  # pylint:disable=g-long-lambda
+              [],
+              minval=1,
+              maxval=10,
+              dtype=dtypes.float32))
+      opt = options_lib.Options()
+      opt.experimental_external_state_policy = (
+          options_lib.ExternalStatePolicy.WARN)
+      dataset0 = dataset0.with_options(opt)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+
+    with ops.device(self._device0):
+      get_next0 = self.getNext(dataset0)
+    with ops.device(self._device1):
+      get_next1 = self.getNext(dataset1)
+    with ops.device(self._device2):
+      get_next2 = self.getNext(dataset2)
+
+    for _ in range(100):
+      self.evaluate(get_next0())
+      self.evaluate(get_next1())
+      self.evaluate(get_next2())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testExternalStatePolicyFail(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: random_ops.random_uniform(  # pylint:disable=g-long-lambda
+              [],
+              minval=1,
+              maxval=10,
+              dtype=dtypes.float32))
+      opt = options_lib.Options()
+      opt.experimental_external_state_policy = (
+          options_lib.ExternalStatePolicy.FAIL)
+      dataset0 = dataset0.with_options(opt)
+    with self.assertRaises(errors.FailedPreconditionError):
+      replicated_ds = distribute.replicate(dataset0,
+                                           [self._device1, self._device2])
+      dataset1 = replicated_ds[self._device1]
+      dataset2 = replicated_ds[self._device2]
+
+      with ops.device(self._device0):
+        get_next0 = self.getNext(dataset0)
+      with ops.device(self._device1):
+        get_next1 = self.getNext(dataset1)
+      with ops.device(self._device2):
+        get_next2 = self.getNext(dataset2)
+
+      for _ in range(100):
+        self.evaluate(get_next0())
+        self.evaluate(get_next1())
+        self.evaluate(get_next2())
+
+
+def _get_server_def(job_name, local_server_port, remote_server_addresses,
+                    task_index):
+  """Returns a server def with a single job + multiple tasks."""
+  cluster_def = cluster_pb2.ClusterDef()
+  job_def = cluster_def.job.add()
+  job_def.name = job_name
+  job_def.tasks[0] = "localhost:%d" % local_server_port
+
+  for i, remote_server_address in enumerate(remote_server_addresses, start=1):
+    job_def.tasks[i] = remote_server_address
+
+  server_def = tensorflow_server_pb2.ServerDef(
+      cluster=cluster_def,
+      job_name=job_name,
+      task_index=task_index,
+      protocol="grpc")
+
+  return server_def
+
+
+class EagerClusterReplicateTest(test_base.DatasetTestBase,
+                                parameterized.TestCase):
+
+  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
+    super(EagerClusterReplicateTest, self).__init__(methodName)
+    self._job_name = "remove_device"
+    self._device0 = "/job:%s/replica:0/task:0/device:CPU:0" % self._job_name
+    self._device1 = "/job:%s/replica:0/task:1/device:CPU:0" % self._job_name
+    self._device2 = "/job:%s/replica:0/task:2/device:CPU:0" % self._job_name
+
+  def setUp(self):
+    super(EagerClusterReplicateTest, self).setUp()
+
+    # TODO(b/171412104): Move create server to __init__ once tfrt support it.
+    self._cached_server1 = server_lib.Server.create_local_server()
+    self._cached_server2 = server_lib.Server.create_local_server()
+    self._cached_server1_target = self._cached_server1.target[len("grpc://"):]
+    self._cached_server2_target = self._cached_server2.target[len("grpc://"):]
+
+    # Start the local server.
+    local_port = pywrap_tfe.TF_PickUnusedPortOrDie()
+    context.set_server_def(
+        server_def=_get_server_def(
+            self._job_name,
+            local_server_port=local_port,
+            remote_server_addresses=[
+                self._cached_server1_target, self._cached_server2_target
+            ],
+            task_index=0))
+
+  def tearDown(self):
+    super().tearDown()
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[2], mode=["eager"]))
+  def testBasic(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    with ops.device(self._device0):
+      self.assertDatasetProduces(dataset0, range(100))
+    with ops.device(self._device1):
+      self.assertDatasetProduces(dataset1, range(100))
+    with ops.device(self._device2):
+      self.assertDatasetProduces(dataset2, range(100))
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[2], mode=["eager"]))
+  def testMap(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100).map(lambda x: x * 2)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    with ops.device(self._device0):
+      self.assertDatasetProduces(dataset0, range(0, 200, 2))
+    with ops.device(self._device1):
+      self.assertDatasetProduces(dataset1, range(0, 200, 2))
+    with ops.device(self._device2):
+      self.assertDatasetProduces(dataset2, range(0, 200, 2))
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[2], mode=["eager"]))
+  def testVariableInput(self):
+    with ops.device(self._device0):
+      counter_var = variable_scope.get_variable(
+          "counter", (), dtypes.int32, use_resource=True)
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: counter_var.assign_add(1))
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    with ops.device(self._device0):
+      self.assertDatasetProduces(
+          dataset0, range(1, 101), requires_initialization=True)
+    with ops.device(self._device1):
+      self.assertDatasetProduces(
+          dataset1, range(1, 101), requires_initialization=True)
+    with ops.device(self._device2):
+      self.assertDatasetProduces(
+          dataset2, range(1, 101), requires_initialization=True)
+
+
+class GraphClusterReplicateTest(test_base.DatasetTestBase,
+                                parameterized.TestCase):
+
+  def setUp(self):
+    super(GraphClusterReplicateTest, self).setUp()
+    # Start the local server.
+    worker_config = config_pb2.ConfigProto()
+    worker_config.device_count["CPU"] = 2
+    worker, _ = test_util.create_local_cluster(
+        3, 0, worker_config=worker_config)
+    self._device0 = "/job:worker/replica:0/task:0/device:CPU:0"
+    self._device1 = "/job:worker/replica:0/task:1/device:CPU:0"
+    self._device2 = "/job:worker/replica:0/task:2/device:CPU:0"
+    self._target = worker[0].target
+
+  def tearDown(self):
+    super().tearDown()
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[1], mode=["graph"]))
+  def testBasic(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    with ops.device(self._device0):
+      get_next = self.getNext(dataset0)
+    with ops.device(self._device1):
+      get_next1 = self.getNext(dataset1)
+    with ops.device(self._device2):
+      get_next2 = self.getNext(dataset2)
+
+    with session.Session(self._target) as sess:
+      for i in range(100):
+        self.assertEqual(i, sess.run(get_next()))
+        self.assertEqual(i, sess.run(get_next1()))
+        self.assertEqual(i, sess.run(get_next2()))
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[1], mode=["graph"]))
+  def testMap(self):
+    with ops.device(self._device0):
+      dataset0 = dataset_ops.Dataset.range(100).map(lambda x: x * 2)
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    dataset2 = replicated_ds[self._device2]
+    with ops.device(self._device0):
+      get_next = self.getNext(dataset0)
+    with ops.device(self._device1):
+      get_next1 = self.getNext(dataset1)
+    with ops.device(self._device2):
+      get_next2 = self.getNext(dataset2)
+
+    with session.Session(self._target) as sess:
+      for i in range(100):
+        self.assertEqual(i * 2, sess.run(get_next()))
+        self.assertEqual(i * 2, sess.run(get_next1()))
+        self.assertEqual(i * 2, sess.run(get_next2()))
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[1], mode=["graph"]))
+  def testVariableInput(self):
+    with ops.device(self._device0):
+      counter_var = variable_scope.get_variable(
+          "counter", (), dtypes.int32, use_resource=True)
+      dataset0 = dataset_ops.Dataset.range(100).map(
+          lambda _: counter_var.assign_add(1))
+    replicated_ds = distribute.replicate(dataset0,
+                                         [self._device1, self._device2])
+    dataset1 = replicated_ds[self._device1]
+    with ops.device(self._device1):
+      it1 = dataset_ops.make_initializable_iterator(dataset1)
+    # We don't support stateful ops across processes in functions as of now.
+    with session.Session(self._target) as sess:
+      with self.assertRaises(errors.OpError):
+        sess.run(it1.initializer)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/BUILD

@@ -0,0 +1,438 @@
+# tf.data service tests.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+tf_py_strict_test(
+    name = "auto_shard_test",
+    size = "medium",
+    srcs = ["auto_shard_test.py"],
+    shard_count = 32,
+    tags = [
+        "no_oss",  # TODO(b/289230931): Investigate flaky test.
+    ],
+    deps = [
+        ":multi_process_cluster",
+        ":test_base",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "coordinated_read_ft_test",
+    size = "medium",
+    srcs = ["coordinated_read_ft_test.py"],
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "coordinated_read_test",
+    size = "medium",
+    srcs = ["coordinated_read_test.py"],
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "data_service_ops_test",
+    size = "medium",
+    srcs = ["data_service_ops_test.py"],
+    shard_count = 32,
+    deps = [
+        ":test_base",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:batching",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:grouping",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/experimental/service:server_lib",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:lookup_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:sparse_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:test",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "gpu_test",
+    srcs = ["gpu_test.py"],
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "dynamic_sharding_test",
+    size = "medium",
+    srcs = ["dynamic_sharding_test.py"],
+    shard_count = 16,
+    deps = [
+        ":test_base",
+        "//tensorflow/core/protobuf:for_core_protos_py_proto",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/service:server_lib",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/kernel_tests:tf_record_test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/logging",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "fault_tolerance_test",
+    size = "small",
+    srcs = ["fault_tolerance_test.py"],
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+        "@pypi//portpicker",
+    ],
+)
+
+tf_py_strict_test(
+    name = "local_workers_test",
+    size = "medium",
+    srcs = ["local_workers_test.py"],
+    shard_count = 24,
+    tags = [
+        "no_oss",  # TODO(b/295501569)
+    ],
+    deps = [
+        ":multi_process_cluster",
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "metadata_test",
+    size = "medium",
+    srcs = ["metadata_test.py"],
+    shard_count = 2,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "cross_trainer_cache_test",
+    size = "medium",
+    srcs = ["cross_trainer_cache_test.py"],
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "cross_trainer_cache_ft_test",
+    size = "medium",
+    srcs = ["cross_trainer_cache_ft_test.py"],
+    shard_count = 4,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "distributed_save_test",
+    size = "medium",
+    srcs = ["distributed_save_test.py"],
+    # copybara:uncomment_begin(google-only)
+    # exec_properties = select({
+    # "//tools/cpp:msan_build": {"mem": "20g"},
+    # "//conditions:default": None,
+    # }),
+    # copybara:uncomment_end
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distributed_save_op",
+        "//tensorflow/python/data/experimental/service:_pywrap_snapshot_utils",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "distributed_save_load_test",
+    size = "medium",
+    srcs = ["distributed_save_load_test.py"],
+    shard_count = 8,
+    tags = [
+        "requires-mem:20g",
+    ],
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distributed_save_op",
+        "//tensorflow/python/data/kernel_tests:checkpoint_test_base",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:test",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "distributed_save_load_ft_test",
+    size = "medium",
+    srcs = ["distributed_save_load_ft_test.py"],
+    # copybara:uncomment_begin(google-only)
+    # exec_properties = select({
+    # "//tools/cpp:msan_build": {"mem": "24g"},
+    # "//conditions:default": {"mem": "20g"},
+    # }),
+    # copybara:uncomment_end
+    shard_count = 8,
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distributed_save_op",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:test",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "distributed_save_ft_test",
+    size = "medium",
+    srcs = ["distributed_save_ft_test.py"],
+    # copybara:uncomment_begin(google-only)
+    # exec_properties = select({
+    # "//tools/cpp:msan_build": {"mem": "20g"},
+    # "//conditions:default": None,
+    # }),
+    # copybara:uncomment_end
+    shard_count = 17,
+    tags = [
+        "no_mac",  # TODO(b/290355883): Fix the flakyness in macos
+    ],
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:distributed_save_op",
+        "//tensorflow/python/data/experimental/service:_pywrap_snapshot_utils",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "multi_device_test",
+    size = "small",
+    srcs = ["multi_device_test.py"],
+    deps = [
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:lookup_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "multi_process_cluster",
+    srcs = ["multi_process_cluster.py"],
+    strict_deps = True,
+    deps = [
+        ":test_base",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/service:server_lib",
+        "//tensorflow/python/distribute:multi_process_lib",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:test",
+    ],
+)
+
+tf_py_strict_test(
+    name = "multi_process_cluster_test",
+    size = "medium",
+    srcs = ["multi_process_cluster_test.py"],
+    shard_count = 3,
+    deps = [
+        ":multi_process_cluster",
+        ":test_base",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_ops",
+        "//tensorflow/python/ops:math_ops",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "test_base",
+    srcs = ["test_base.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/experimental/service:server_lib",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:test",
+    ],
+)
+
+tf_py_strict_test(
+    name = "worker_tags_test",
+    size = "medium",
+    srcs = ["worker_tags_test.py"],
+    shard_count = 32,
+    tags = [
+        "no_oss_py313",  # TODO(b/427743382): Fix the timeout in linux_x86_cpu_wheel_py313_np1
+    ],
+    deps = [
+        ":multi_process_cluster",
+        ":test_base",
+        "//tensorflow/python/data/experimental/ops:data_service_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/data/experimental/kernel_tests/service/auto_shard_test.py

@@ -0,0 +1,575 @@
+# Copyright 2021 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 auto-sharding datasets with tf.data service."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import data_service_pb2
+from tensorflow.python.data.experimental.kernel_tests.service import multi_process_cluster
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.kernel_tests.service.multi_process_cluster import MultiProcessCluster
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distribute
+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 readers
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+
+
+def _make_service_cluster(
+    num_workers,
+    local_shard_index,
+    worker_addresses=None,
+    deployment_mode=data_service_pb2.DEPLOYMENT_MODE_COLOCATED):
+  if worker_addresses is None:
+    worker_addresses = ["localhost" for _ in range(num_workers)]
+
+  cluster = MultiProcessCluster(
+      num_local_workers=0,
+      num_remote_workers=0,
+      worker_addresses=worker_addresses,
+      deployment_mode=deployment_mode)
+  for _ in range(local_shard_index):
+    cluster.start_remote_worker()
+  cluster.start_local_worker()
+  for _ in range(num_workers - local_shard_index - 1):
+    cluster.start_remote_worker()
+  return cluster
+
+
+# pylint:disable=g-complex-comprehension
+class AutoShardTest(data_service_test_base.TestBase,
+                    tf_record_test_base.TFRecordTestBase,
+                    parameterized.TestCase):
+  """Tests auto-sharding datasets with tf.data service."""
+
+  def setUp(self):
+    super(AutoShardTest, self).setUp()
+    self._num_files = 10
+    self._num_records = 10
+    self._filenames = self._createFiles()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.DATA,
+              data_service_ops.ShardingPolicy.FILE_OR_DATA
+          ])))
+  def testRangeDataset_AutoShard(self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+    self.assertDatasetProduces(dataset, [1, 6, 11, 16])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_FileShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE)
+    with self.assertRaisesRegex(errors.NotFoundError,
+                                "Found an unshardable source dataset"):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(worker_index=[distribute.SHARD_HINT, 0, 5])))
+  def testRangeDataset_ShardHint(self, worker_index):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    # With HINT sharding, `num_shards` should be `SHARD_HINT`; `index` can be
+    # any value.
+    dataset = dataset.shard(
+        num_shards=distribute.SHARD_HINT, index=worker_index)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT)
+    self.assertDatasetProduces(dataset, [1, 6, 11, 16])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_InvalidWorkerIndexUsingShardHint(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    # With HINT sharding, `SHARD_HINT` should be passed to `num_shards`, not
+    # `index`.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        r"Index must be between 0 and 4 \(currently index = -1\)."):
+      dataset = dataset.shard(num_shards=5, index=distribute.SHARD_HINT)
+      dataset = self.make_distributed_dataset(
+          dataset,
+          cluster=cluster,
+          processing_mode=data_service_ops.ShardingPolicy.HINT)
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_NoShardHint(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    # No SHARD_HINT is provided. The given sharding arguments will be used.
+    dataset = dataset.shard(num_shards=1, index=0)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT)
+    self.assertDatasetProduces(dataset, list(range(20)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.OFF,
+              data_service_ops.ShardingPolicy.FILE_OR_DATA
+          ])))
+  def testRangeDataset_ShardHintUsedInWrongShardingPolicy(
+      self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = dataset.shard(distribute.SHARD_HINT, distribute.SHARD_HINT)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "tf.data service with "
+        "`tf.data.experimental.service.ShardingPolicy.HINT` processing mode."):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_NoShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        target_workers="LOCAL")
+    self.assertDatasetProduces(dataset, list(range(20)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_OneWorker(self):
+    """Makes sure shards from all workers form the complete dataset."""
+    cluster = _make_service_cluster(num_workers=1, local_shard_index=0)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    self.assertDatasetProduces(dataset, list(range(20)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRangeDataset_ReadFromAnyWorker(self):
+    # When deployment mode is unspecified, the client will read from any worker.
+    cluster = _make_service_cluster(
+        num_workers=5, local_shard_index=1, deployment_mode=None)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    self.assertDatasetProduces(
+        dataset, list(range(20)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.FILE_OR_DATA,
+              data_service_ops.ShardingPolicy.FILE
+          ])))
+  def testTFRecordDataset_AutoShard(self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=sharding_policy,
+        target_workers="LOCAL")
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in (3, 8)
+        for record in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.FILE_OR_DATA,
+              data_service_ops.ShardingPolicy.FILE
+          ])))
+  def testTFRecordDataset_ShuffleFileList(self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=True)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in (3, 8)
+        for record in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_DataShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DATA)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 10)
+        for record in (3, 8)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_HintDataShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = dataset.shard(distribute.SHARD_HINT, distribute.SHARD_HINT)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 10)
+        for record in (3, 8)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_HintFileShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.shard(distribute.SHARD_HINT, distribute.SHARD_HINT)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in (3, 8)
+        for record in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_NoShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        target_workers="LOCAL")
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 10)
+        for record in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_ReadFromAnyWorker(self):
+    # When deployment mode is unspecified, the client will read from any worker.
+    cluster = _make_service_cluster(
+        num_workers=5, local_shard_index=3, deployment_mode=None)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 10)
+        for record in range(0, 10)
+    ]
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.FILE_OR_DATA,
+              data_service_ops.ShardingPolicy.FILE
+          ])))
+  def testTFRecordDataset_FewerFilesThanWorkers(self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames[:4], shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "not enough for the required 5 shards/workers."):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_FewerFilesThanWorkers_HintShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames[:4], shuffle=False)
+    dataset = dataset.shard(distribute.SHARD_HINT, distribute.SHARD_HINT)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT)
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "not enough for the required 5 shards/workers."):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTFRecordDataset_FewerFilesThanWorkers_DataShard(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames[:4], shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DATA)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 4)
+        for record in (3, 8)
+    ]
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.FILE_OR_DATA,
+              data_service_ops.ShardingPolicy.DATA
+          ])))
+  def testBatchDataset(self, sharding_policy):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=1)
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = dataset.batch(batch_size=3, drop_remainder=False)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+    self.assertDatasetProduces(dataset, [[3, 4, 5], [18, 19]])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInterleaveDataset(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.interleave(
+        readers.TFRecordDataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for record in range(0, 10)
+        for file in (3, 8)
+    ]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipDataset(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset1 = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset1 = dataset1.interleave(
+        readers.TFRecordDataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset2 = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset2 = dataset2.interleave(
+        readers.TFRecordDataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+
+    expected = [(b"Record %d of file %d" % (record, file),
+                 b"Record %d of file %d" % (record, file))
+                for record in range(0, 10)
+                for file in (3, 8)]
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcatenateDataset(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset1 = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset1 = dataset1.interleave(
+        readers.TFRecordDataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset2 = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset2 = dataset2.interleave(
+        readers.TFRecordDataset,
+        cycle_length=10,
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset = dataset1.concatenate(dataset2)
+    dataset = dataset.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for record in range(0, 10)
+        for file in (3, 8)
+    ]
+    expected += expected
+    self.assertDatasetProduces(dataset, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyDataset(self):
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.range(0)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    self.assertDatasetProduces(dataset, [])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAnonymousPorts(self):
+    cluster = _make_service_cluster(
+        num_workers=5,
+        local_shard_index=3,
+        worker_addresses=["localhost:%port%" for _ in range(5)])
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    self.assertDatasetProduces(dataset, [3, 8, 13, 18])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNamedPorts(self):
+    cluster = _make_service_cluster(
+        num_workers=5,
+        local_shard_index=3,
+        worker_addresses=["localhost:%port_worker%" for _ in range(5)])
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    self.assertDatasetProduces(dataset, [3, 8, 13, 18])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidPorts(self):
+    with self.assertRaisesRegex(RuntimeError,
+                                "The worker's address is not configured"):
+      _ = _make_service_cluster(
+          num_workers=5,
+          local_shard_index=0,
+          worker_addresses=["localhost:worker" for _ in range(5)])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyWorkerList(self):
+    cluster = _make_service_cluster(
+        num_workers=5, local_shard_index=1, worker_addresses=[])
+    dataset = dataset_ops.Dataset.range(20)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    with self.assertRaisesRegex(errors.NotFoundError,
+                                "Worker .* is not in the workers list."):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWorkerNotFound(self):
+    worker_addresses = [f"fake_worker_{i}" for i in range(5)]
+    with self.assertRaisesRegex(RuntimeError,
+                                "The worker's address is not configured"):
+      _ = _make_service_cluster(
+          num_workers=5, local_shard_index=0, worker_addresses=worker_addresses)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMoreWorkersThanConfigured(self):
+    worker_addresses = ["localhost:%port%"]
+    with self.assertRaisesRegex(
+        RuntimeError,
+        "other workers are already running at the configured host"):
+      _ = _make_service_cluster(
+          num_workers=5, local_shard_index=1, worker_addresses=worker_addresses)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoLocalWorkers(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=0, num_remote_workers=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Static sharding policy <FILE_OR_DATA> requires local tf.data workers"):
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              sharding_policy=list(data_service_ops.ShardingPolicy))))
+  def testEnumerateShardingPolicies(self, sharding_policy):
+    """Verifies tf.data service handles every sharding policy with no errors."""
+    cluster = _make_service_cluster(num_workers=5, local_shard_index=3)
+    dataset = dataset_ops.Dataset.list_files(self._filenames, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+    self.getDatasetOutput(dataset)
+
+
+if __name__ == "__main__":
+  multi_process_cluster.test_main()

tensorflow/python/data/experimental/kernel_tests/service/coordinated_read_ft_test.py

@@ -0,0 +1,121 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Fault tolerance tests for tf.data service coordinated reads."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+class CoordinatedReadFTTest(data_service_test_base.TestBase,
+                            parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(workers_to_add=[1, 3, 10])))
+  def testAddWorkers(self, workers_to_add):
+    starting_workers = 3
+    cluster = data_service_test_base.TestCluster(num_workers=starting_workers)
+    num_consumers = 7
+    ds = self.make_coordinated_read_dataset(cluster, num_consumers)
+
+    get_next = self.getNext(ds, requires_initialization=True)
+    results = []
+    zeros_seen = 0
+    for _ in range(25):
+      results.append(self.evaluate(get_next()))
+      if results[-1] == 0:
+        zeros_seen += 1
+    for _ in range(workers_to_add):
+      cluster.add_worker()
+    # Read until all new workers have joined.
+    while zeros_seen < starting_workers + workers_to_add:
+      results.append(self.evaluate(get_next()))
+      if results[-1] == 0:
+        zeros_seen += 1
+    # Read some more.
+    for _ in range(25):
+      results.append(self.evaluate(get_next()))
+
+    self.checkCoordinatedReadGroups(results, num_consumers)
+    cluster.stop_workers()
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRestartWorker(self):
+    num_workers = 3
+    # Set a shutdown quiet period to prevent workers from shutting down partway
+    # through a round.
+    cluster = data_service_test_base.TestCluster(
+        num_workers, worker_shutdown_quiet_period_ms=2000)
+    num_consumers = 5
+    ds = self.make_coordinated_read_dataset(cluster, num_consumers)
+
+    get_next = self.getNext(ds, requires_initialization=True)
+    results = []
+
+    self.read(get_next, results, 20)
+    cluster.workers[1].stop()
+    # Check that we can continue to read even with a worker stopped.
+    self.read(get_next, results, 20)
+    cluster.workers[1].restart()
+    # Read until we get results from the restarted worker, then read some more.
+    while results[-1] != 0:
+      results.append(self.evaluate(get_next()))
+    self.read(get_next, results, 20)
+    self.checkCoordinatedReadGroups(results, num_consumers)
+    cluster.stop_workers()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(sharding_policy=[
+              data_service_ops.ShardingPolicy.OFF,
+              data_service_ops.ShardingPolicy.DYNAMIC
+          ])))
+  def testMultiStartStop(self, sharding_policy):
+    num_workers = 3
+    # Set a shutdown quiet period to prevent workers from shutting down partway
+    # through a round.
+    cluster = data_service_test_base.TestCluster(
+        num_workers, worker_shutdown_quiet_period_ms=2000)
+    num_consumers = 5
+    ds = self.make_coordinated_read_dataset(cluster, num_consumers,
+                                            sharding_policy)
+
+    get_next = self.getNext(ds, requires_initialization=True)
+    results = []
+
+    self.read(get_next, results, 20)
+    for i in range(num_workers):
+      cluster.workers[i].stop()
+      self.read(get_next, results, 20)
+      cluster.workers[i].restart()
+      self.read(get_next, results, 20)
+
+    cluster.add_worker()
+    cluster.restart_dispatcher()
+    for i in range(num_workers):
+      cluster.workers[i].stop()
+    self.read(get_next, results, 20)
+
+    self.checkCoordinatedReadGroups(results, num_consumers)
+    cluster.stop_workers()
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/coordinated_read_test.py

@@ -0,0 +1,188 @@
+# Copyright 2021 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 service coordinated reads."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import grouping
+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.platform import test
+
+
+class CoordinatedReadTest(data_service_test_base.TestBase,
+                          parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3], num_consumers=[1, 2, 5])))
+  def testBasic(self, num_workers, num_consumers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    ds = self.make_coordinated_read_dataset(cluster, num_consumers)
+    get_next = self.getNext(ds, requires_initialization=True)
+    results = [self.evaluate(get_next()) for _ in range(100)]
+    self.checkCoordinatedReadGroups(results, num_consumers)
+    cluster.stop_workers()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConsumerRestart(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_consumers = 3
+    ds = self.make_coordinated_read_dataset(cluster, num_consumers)
+    get_next = self.getNext(ds, requires_initialization=True)
+    _ = [self.evaluate(get_next()) for _ in range(20)]
+
+    ds2 = self.make_coordinated_read_dataset(cluster, num_consumers)
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                "current round has already reached"):
+      get_next_ds2 = self.getNext(ds2, requires_initialization=True)
+      _ = [self.evaluate(get_next_ds2()) for _ in range(20)]
+    cluster.stop_workers()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBucketizing(self):
+    # Tests a common use case for round robin reads. At each step, all
+    # consumers should get batches with the same bucket size.
+    cluster = data_service_test_base.TestCluster(num_workers=4)
+    num_elements = 100
+    low_bucket_max = 30
+    mid_bucket_max = 60
+    bucket_boundaries = [low_bucket_max, mid_bucket_max]
+    batch_size = 10
+    num_consumer_hosts = 3
+    replicas_per_consumer_host = 5
+    num_consumers = num_consumer_hosts * replicas_per_consumer_host
+    bucket_batch_sizes = [batch_size] * (len(bucket_boundaries) + 1)
+    # Set up the dataset that will run on the tf.data workers.
+    ds = dataset_ops.Dataset.range(num_elements, output_type=dtypes.int32)
+    ds = ds.shuffle(num_elements)
+    ds = ds.repeat()
+    ds = ds.apply(
+        grouping.bucket_by_sequence_length(
+            lambda x: x,
+            bucket_boundaries,
+            bucket_batch_sizes,
+            drop_remainder=True))
+    ds = ds.apply(
+        grouping.group_by_window(
+            lambda x: math_ops.cast(x[1], dtypes.int64),
+            lambda _, x: dataset_ops.Dataset.from_tensors(x),
+            window_size=num_consumers))
+    ds = ds.flat_map(lambda x: x)
+
+    # Set up the per-consumer-host datasets. During each global step, we pull
+    # `replicas_per_consumer_host` batches from each of these datasets.
+    host_datasets = []
+    for host_index in range(num_consumer_hosts):
+      per_replica_datasets = []
+      for i in range(replicas_per_consumer_host):
+        consumer_index = host_index * replicas_per_consumer_host + i
+        per_replica_datasets.append(
+            self.make_distributed_dataset(
+                ds,
+                cluster,
+                job_name="test",
+                consumer_index=consumer_index,
+                num_consumers=num_consumers))
+      host_dataset = dataset_ops.Dataset.from_tensor_slices(
+          per_replica_datasets)
+      host_dataset = host_dataset.interleave(
+          lambda x: x,
+          cycle_length=len(per_replica_datasets),
+          num_parallel_calls=len(per_replica_datasets),
+          deterministic=True)
+      host_datasets.append(host_dataset)
+
+    # Use parallel interleave to read from host datasets in parallel.
+    ds = dataset_ops.Dataset.from_tensor_slices(host_datasets)
+    ds = ds.interleave(
+        lambda x: x,
+        block_length=replicas_per_consumer_host,
+        cycle_length=len(host_datasets),
+        num_parallel_calls=len(host_datasets),
+        deterministic=True)
+
+    num_rounds = 4
+    get_next = self.getNext(ds, requires_initialization=True)
+    results = []
+    for i in range(num_rounds * num_consumers):
+      results.append(self.evaluate(get_next()))
+
+    def get_bucket(elem):
+      bucket_ind = 0
+      while bucket_ind < len(
+          bucket_boundaries) and elem >= bucket_boundaries[bucket_ind]:
+        bucket_ind += 1
+      return bucket_ind
+
+    # Check that the batches for each step contain elements from the same
+    # bucket.
+    for i in range(0, len(results), num_consumers):
+      batches = results[num_consumers * i:num_consumers * (i + 1)]
+      bucket_inds = [get_bucket(batch[0]) for batch in batches]
+      for bucket_ind in bucket_inds[1:]:
+        self.assertEqual(
+            bucket_inds[0], bucket_ind,
+            "Batches: {}, Buckets: {}".format(batches, bucket_inds))
+
+  @combinations.generate(test_base.v1_only_combinations())
+  def testFiniteV1(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self.make_distributed_dataset(
+        ds, cluster, job_name="test", consumer_index=0, num_consumers=1)
+
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "Encountered end of sequence on a "
+        "round-robin read iterator"):
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.v2_only_combinations())
+  def testFiniteV2(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self.make_distributed_dataset(
+        ds, cluster, job_name="test", consumer_index=0, num_consumers=1)
+
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "Round robin reads "
+        "require that the input dataset has infinite "
+        "cardinality, but the dataset has cardinality " + str(num_elements)):
+      self.getDatasetOutput(ds)
+
+  # We test only eager combinations because the `map` transformation used for
+  # compression in make_distributed_dataset makes cardinality unknown in TF1.
+  @combinations.generate(test_base.v2_only_combinations())
+  def testCardinality(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    ds = self.make_distributed_dataset(
+        dataset_ops.Dataset.range(10).repeat(),
+        cluster,
+        job_name="test",
+        consumer_index=0,
+        num_consumers=2)
+    self.assertEqual(self.evaluate(ds.cardinality()), dataset_ops.INFINITE)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_ft_test.py

@@ -0,0 +1,94 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Fault tolerance tests for tf.data service cross-trainer cache."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+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 CrossTrainerCacheFtTest(data_service_test_base.TestBase,
+                              parameterized.TestCase):
+  """Fault tolerance tests for tf.data service cross-trainer cache."""
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWorkerRestart(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    distributed_dataset = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+
+    get_next = self.getNext(distributed_dataset)
+    elements = self._get_next(get_next, 100)
+    self.assertEqual(elements, list(range(100)))
+
+    cluster.workers[0].restart()
+
+    # Read until we get results from the restarted worker, then read some more.
+    while self.evaluate(get_next()) != 0:
+      pass
+
+    elements = self._get_next(get_next, 100)
+    self.assertEqual(elements, list(range(1, 101)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDispatcherRestart(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    distributed_dataset = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+
+    get_next = self.getNext(distributed_dataset)
+    elements = self._get_next(get_next, 100)
+    self.assertEqual(elements, list(range(100)))
+
+    cluster.restart_dispatcher()
+
+    # Dispatcher restart should not affect the workers.
+    elements = self._get_next(get_next, 100)
+    self.assertEqual(elements, list(range(100, 200)))
+
+  def _get_next(self, get_next, num_elements):
+    return [self.evaluate(get_next()) for _ in range(num_elements)]
+
+  def _create_cluster(self,
+                      num_workers,
+                      cross_trainer_cache_size_bytes=10 * (2**30)):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    for _ in range(num_workers):
+      worker = data_service_test_base.TestWorker(
+          dispatcher_address=cluster.dispatcher_address(),
+          shutdown_quiet_period_ms=0,
+          cross_trainer_cache_size_bytes=cross_trainer_cache_size_bytes)
+      worker.start()
+      cluster.workers.append(worker)
+    return cluster
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py

@@ -0,0 +1,464 @@
+# Copyright 2022 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 sharing datasets across training jobs."""
+
+import multiprocessing
+
+from absl.testing import parameterized
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+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.platform import test
+
+
+class CrossTrainerCacheTest(data_service_test_base.TestBase,
+                            parameterized.TestCase):
+  """Tests for sharing datasets across jobs using a cross-trainer cache."""
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEnableCrossTrainerCache(self):
+    """Tests cross-trainer cache with `distribute`."""
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    # The second client reads the same data from the cross-trainer cache.
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    self.assertDatasetProduces(dataset2.take(10), list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFromDatasetId(self):
+    """Tests cross-trainer cache with `register_dataset`/`from_dataset_id`."""
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset_id1 = data_service_ops.register_dataset(
+        cluster.dispatcher.target, dataset, dataset_id="dataset_id")
+    dataset1 = data_service_ops.from_dataset_id(
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        service=cluster.dispatcher.target,
+        dataset_id=dataset_id1,
+        element_spec=dataset.element_spec,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    dataset_id2 = data_service_ops.register_dataset(
+        cluster.dispatcher.target, dataset, dataset_id="dataset_id")
+    dataset2 = data_service_ops.from_dataset_id(
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        service=cluster.dispatcher.target,
+        dataset_id=dataset_id2,
+        element_spec=dataset.element_spec,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    self.assertDatasetProduces(dataset2.take(10), list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDisableCrossTrainerCacheByDefault(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(dataset, cluster, job_name="job")
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    # The two clients use the same job. The second client can't read the data
+    # already read by the first client.
+    dataset2 = self.make_distributed_dataset(dataset, cluster, job_name="job")
+    output = self.getDatasetOutput(dataset2.take(10))
+    self.assertGreaterEqual(output[0], 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testConcurrentReaders(self):
+    # Fetching an element from the dataset will trigger prefetches of more
+    # elements, one per CPU core which will be placed in the cache.
+    # However if the number of prefetches exceeds the space available in
+    # the cache then the sliding window will be moved forward away from
+    # the element just read thus negating the use of the cache as other
+    # trainers will not get the correct element.
+    # Hence the need to calculate the size of the cache based on the
+    # number of CPU cores and the element size of 423. The extra 8
+    # entries are simply a bit of margin.
+    num_cpus = multiprocessing.cpu_count()
+    cluster = self._create_cluster(
+        num_workers=1, cross_trainer_cache_size_bytes=(num_cpus + 8) * 423)
+    num_readers = 20
+    num_elements = 50
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+
+    datasets = []
+    iterators = []
+    for i in range(num_readers):
+      distributed_dataset = self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name="job",
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id=f"Trainer {i}"),
+          max_outstanding_requests=1)
+      iterator = self.getNext(distributed_dataset)
+      datasets.append(distributed_dataset)
+      iterators.append(iterator)
+
+    for i in range(num_elements):
+      # All the readers read the same element in one step.
+      for j in range(num_readers):
+        self.assertEqual(self.evaluate(iterators[j]()), i)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSlowClientSkipsData(self):
+    cluster = self._create_cluster(
+        num_workers=1, cross_trainer_cache_size_bytes=500)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(200), list(range(200)))
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    dataset2 = dataset2.take(200)
+    output = self.getDatasetOutput(dataset2)
+    # When the cache is small, the second trainer couldn't read the beginning of
+    # the dataset. It can still read 200 elements from the dataset, because the
+    # dataset is infinite.
+    self.assertGreater(output[0], 0)
+    self.assertLen(output, 200)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSmallCache(self):
+    cluster = self._create_cluster(
+        num_workers=1, cross_trainer_cache_size_bytes=500)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    num_readers = 20
+
+    for i in range(num_readers):
+      # Even if the cache is small and may discard old data, each trainer can
+      # still read the required number of elements because the input dataset is
+      # infinite.
+      distributed_dataset = self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name="job",
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id=f"Trainer {i}"))
+      output = self.getDatasetOutput(distributed_dataset.take(200))
+      self.assertLen(output, 200)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleDataset(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat().shuffle(
+        buffer_size=100)
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    output1 = self.getDatasetOutput(dataset1.take(10))
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    output2 = self.getDatasetOutput(dataset2.take(10))
+    self.assertEqual(output1, output2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSameTrainerID(self):
+    # Jobs from the same training cluster do not reuse data from the cache.
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer ID"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer ID"))
+    output = self.getDatasetOutput(dataset2.take(10))
+    self.assertGreaterEqual(output[0], 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentJobNames(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job1",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job2",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    self.assertDatasetProduces(dataset2.take(10), list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDynamicSharding(self):
+    cluster = self._create_cluster(num_workers=2)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    output1 = self.getDatasetOutput(dataset1.take(100))
+
+    # The second client reads the same data from the cross-trainer cache.
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    output2 = self.getDatasetOutput(dataset2.take(100))
+    # Verifies the intersection is non-empty.
+    self.assertTrue(set(output1) & set(output2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoCompression(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        compression=None,
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        compression=None,
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    self.assertDatasetProduces(dataset2.take(10), list(range(10)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCompressionMismatch(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Data type mismatch"):
+      dataset2 = self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name="job",
+          compression=None,
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id="Trainer 1"))
+      self.getDatasetOutput(dataset2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRequiresJobName(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Cross-trainer caching requires named jobs. Got empty `job_name`."):
+      dataset = self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name=None,
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id="Trainer 1"))
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(range_=[0, 10])))
+  def testRequiresInfiniteDataset(self, range_):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(range_).map(lambda x: x + 1)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Cross-trainer caching requires the input dataset to be infinite."):
+      dataset = dataset.apply(
+          data_service_ops.distribute(
+              processing_mode=data_service_ops.ShardingPolicy.OFF,
+              service=cluster.dispatcher.target,
+              job_name="job_name",
+              cross_trainer_cache=data_service_ops.CrossTrainerCache(
+                  trainer_id="Trainer ID")))
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testMultipleIterationsForOneDatasetEagerMode(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    # In the eager mode, each iteration creates a new data service job and does
+    # not reuse cached data. We disallow this use case.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Cross-trainer caching requires infinite datasets and disallows "
+        "multiple repetitions of the same dataset."):
+      self.getDatasetOutput(dataset1.take(10))
+      self.getDatasetOutput(dataset1.take(10))
+      self.getDatasetOutput(dataset1.take(10))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testMultipleIterationsForOneDatasetGraphMode(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    # These clients are assumed to be from the same training cluster. Thus, they
+    # do not reuse data from the cross-trainer cache.
+    output1 = self.getDatasetOutput(dataset1.take(10))
+    output1 += self.getDatasetOutput(dataset1.take(10))
+    output1 += self.getDatasetOutput(dataset1.take(10))
+    self.assertLen(set(output1), 30)
+
+    dataset2 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 2"))
+    # These clients reuse some data from the previous clients (not exactly the
+    # same data due to client-side buffering).
+    output2 = self.getDatasetOutput(dataset2.take(10))
+    output2 += self.getDatasetOutput(dataset2.take(10))
+    output2 += self.getDatasetOutput(dataset2.take(10))
+    self.assertTrue(set(output1) & set(output2))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDisallowCoordinatedRead(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Cross-trainer caching does not support coordinated reads."):
+      dataset = self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name="job",
+          num_consumers=1,
+          consumer_index=0,
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id="Trainer 1"))
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNamedJobMismatch(self):
+    cluster = self._create_cluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+    dataset1 = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        job_name="job",
+        cross_trainer_cache=data_service_ops.CrossTrainerCache(
+            trainer_id="Trainer 1"))
+    self.assertDatasetProduces(dataset1.take(10), list(range(10)))
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Existing cross-trainer cache: <enabled>; got <disabled>"):
+      dataset2 = self.make_distributed_dataset(
+          dataset, cluster, job_name="job", cross_trainer_cache=None)
+      self.getDatasetOutput(dataset2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRequiresNonEmptyTrainerID(self):
+    cluster = self._create_cluster(num_workers=2)
+    dataset = dataset_ops.Dataset.range(10000000).repeat()
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "tf.data service cross-trainer cache requires a non-empty trainer ID."):
+      self.make_distributed_dataset(
+          dataset,
+          cluster,
+          job_name="job",
+          cross_trainer_cache=data_service_ops.CrossTrainerCache(
+              trainer_id=None))
+
+  def _create_cluster(self,
+                      num_workers,
+                      cross_trainer_cache_size_bytes=10 * (2**30)):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    for _ in range(num_workers):
+      worker = data_service_test_base.TestWorker(
+          dispatcher_address=cluster.dispatcher_address(),
+          shutdown_quiet_period_ms=0,
+          cross_trainer_cache_size_bytes=cross_trainer_cache_size_bytes)
+      worker.start()
+      cluster.workers.append(worker)
+    return cluster
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py

@@ -0,0 +1,536 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Fault tolerance tests for tf.data service snapshots."""
+
+import collections
+import os
+import pathlib
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import distributed_save_op
+from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
+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.platform import test
+
+
+def write_file(path):
+  os.makedirs(os.path.dirname(path), exist_ok=True)
+  with open(path, "w") as _:
+    pass
+
+
+def get_stream_assignment(
+    cluster,
+    worker_idx,
+    path,
+    block=True,
+    active_only=False):
+  while True:
+    for progress in cluster.workers[worker_idx].snapshot_task_progresses():
+      if (progress.snapshot_task_base_path.decode() == path
+          and not (active_only and progress.completed)):
+        return progress.snapshot_task_stream_index
+    if not block:
+      break
+    time.sleep(0.1)
+
+
+def get_stream_assignments(
+    cluster,
+    num_workers,
+    paths,
+    block=True,
+    active_only=False):
+  assignments = collections.defaultdict(dict)
+  for worker_idx in range(num_workers):
+    for path in paths:
+      assignment = get_stream_assignment(
+          cluster, worker_idx, path, block, active_only)
+      if assignment is not None:
+        assignments[worker_idx][path] = assignment
+  return assignments
+
+
+def snapshot_is_done(path):
+  return os.path.exists(
+      _pywrap_snapshot_utils.TF_DATA_SnapshotDoneFilePath(path))
+
+
+def snapshot_has_error(path):
+  return os.path.exists(
+      _pywrap_snapshot_utils.TF_DATA_SnapshotErrorFilePath(path))
+
+
+def snapshots_are_done(paths):
+  return all([snapshot_is_done(path) for path in paths])
+
+
+def wait_for_snapshot(paths, f=lambda: None):
+  if isinstance(paths, str):
+    paths = [paths]
+  while not all([snapshot_is_done(path) or snapshot_has_error(path)
+                 for path in paths]):
+    f()
+    time.sleep(0.1)
+
+
+class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
+
+  maxDiff = None
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoverySucceeds(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = self._get_dataset()
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryBlocksOverwrite(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = self._get_dataset()
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    cluster.restart_dispatcher()
+    with self.assertRaisesRegex(
+        errors.AlreadyExistsError, "is already started or completed"):
+      self.evaluate(distributed_save_op.distributed_save(
+          dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+  # TODO(b/250921378): Figure out why tsan times out when there is a worker.
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              bad_stream_dir_name=["stream_", "stream_x", "stream_-1"])))
+  def testSnapshotRecoveryFailsWithBadStreamName(self, bad_stream_dir_name):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    self._make_stream_dir(snapshot_dir.full_path, bad_stream_dir_name)
+    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              bad_source_dir_name=["source_", "source_x", "source_-1"])))
+  def testSnapshotRecoveryFailsWithBadSourceName(self, bad_source_dir_name):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
+                             bad_source_dir_name))
+    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryFailsWithOutOfBoundsSourceName(self):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
+                             "source_1"))
+    with self.assertRaisesRegex(RuntimeError, "Found conflict"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              bad_split_filename=[
+                  "split_",
+                  "split_x_0",
+                  "split_-1_0",
+                  "split_0_x",
+                  "split_0_-1"])))
+  def testSnapshotRecoveryFailsWithBadSplitNames(self, bad_split_filename):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
+                            bad_split_filename))
+    with self.assertRaisesRegex(
+        ValueError,
+        "Expected split_<local_split_index>_<global_split_index>"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryFailsWithOutOfOrderSplitName(self):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
+                            "split_1_0"))
+    with self.assertRaisesRegex(
+        ValueError,
+        "The local split index 1 exceeds the global split index 0"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryFailsWithMissingGlobalIndexInSplitNames(self):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
+                            "split_0_1"))
+    with self.assertRaisesRegex(RuntimeError, "Found missing global"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryFailsWithDuplicateGlobalIndexInSplitName(self):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    write_file(os.path.join(self._source_dir(
+        snapshot_dir.full_path, stream_idx=0), "split_0_1"))
+    write_file(os.path.join(self._source_dir(
+        snapshot_dir.full_path, stream_idx=1, worker=1), "split_0_1"))
+    with self.assertRaisesRegex(RuntimeError, "Found duplicate global"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotRecoveryFailsWithDuplicateWorkerAssignment(self):
+    cluster = data_service_test_base.TestCluster(num_workers=0)
+    snapshot_dir = data_service_test_base.TempDir()
+    self.evaluate(distributed_save_op.distributed_save(
+        self._get_dataset(),
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+
+    write_file(os.path.join(
+        self._source_dir(snapshot_dir.full_path, stream_idx=0), "split_0_1"))
+    write_file(os.path.join(
+        self._source_dir(snapshot_dir.full_path, stream_idx=1), "split_0_1"))
+    with self.assertRaisesRegex(RuntimeError, "worker is already assigned"):
+      cluster.restart_dispatcher()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStreamsReassignedAfterDispatcherRestart(self):
+    n = 5
+    cluster = data_service_test_base.TestCluster(num_workers=n)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = self._get_dataset(dataset_range=10000)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    get_streams = lambda: cluster.snapshot_streams(snapshot_dir.full_path)
+    while len(get_streams()) != n:
+      time.sleep(0.1)
+    cluster.restart_dispatcher()
+    streams = get_streams()
+    while len(streams) != n:
+      time.sleep(0.1)
+      streams = get_streams()
+    self.assertCountEqual([stream.index for stream in streams], range(n))
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              worker_max_concurrent_snapshots=[1, 2])))
+  def testWorkersDontExceedMaxStreamAssignments(
+      self, worker_max_concurrent_snapshots):
+    num_workers = 2
+    num_snapshots = 10
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers,
+        worker_max_concurrent_snapshots=worker_max_concurrent_snapshots)
+    snapshot_dir = data_service_test_base.TempDir()
+    paths = []
+    for i in range(num_snapshots):
+      paths.append(f"{snapshot_dir.full_path}_{i}")
+      self.evaluate(
+          distributed_save_op.distributed_save(
+              dataset_ops.Dataset.range(5000),
+              paths[i],
+              cluster.dispatcher_address()))
+
+    # A mapping of worker idx to max active assignments observed at any time.
+    max_assignments = collections.defaultdict(int)
+
+    def get_assignments_and_update_max_assignments():
+      assignments = get_stream_assignments(
+          cluster, num_workers, paths, block=False, active_only=True)
+      for worker_idx, worker_assignments in assignments.items():
+        max_assignments[worker_idx] = max(max_assignments[worker_idx],
+                                          len(worker_assignments))
+      return assignments
+
+    # Blocks until each worker has at least the max expected active assignments.
+    while True:
+      assignments = get_assignments_and_update_max_assignments()
+      all_workers_have_assignments = len(assignments) == num_workers
+      each_worker_has_enough_assignments = all([
+          len(per_worker_assignments) >= worker_max_concurrent_snapshots
+          for per_worker_assignments in assignments.values()])
+      if all_workers_have_assignments and each_worker_has_enough_assignments:
+        break
+      time.sleep(0.1)
+
+    cluster.restart_dispatcher()
+    wait_for_snapshot(paths, get_assignments_and_update_max_assignments)
+    self.assertValuesEqual(list(max_assignments.values()),
+                           [worker_max_concurrent_snapshots] * num_workers)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def testDatasetRecoversAndCompletes(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(1000)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset,
+            snapshot_dir.full_path,
+            cluster.dispatcher_address(),
+            compression=None))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(dataset, range(1000), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3], num_sources=[1, 3])))
+  def testLargeMultiSourceSnapshotRecoversAndCompletes(
+      self, num_workers, num_sources):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = self._get_dataset(dataset_range=1000, num_sources=num_sources)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+
+    if num_workers > 1:
+      # Can't verify the elements with more than 1 worker since the splits
+      # receive by each worker for each source is non-deterministic.
+      # For example, with dynamic sharding, worker 1 may receive 999 splits for
+      # source 0 and 0 splis for source 1; worker 2 may receive 0 splits for
+      # source 0 and 999 splits for source 1. In this case, the snapshot will
+      # contain no elements since `zip` drops the elements when the zipped
+      # datasets have different cardinalities.
+      return
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    expected = (
+        list(range(1000))
+        if num_sources == 1
+        else list(zip(*([range(1000)] * num_sources))))
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_elements=[1, 2, 1000],
+              num_workers=[1, 3],
+              num_repetitions=[1, 10])))
+  def testRepeatedDatasetRecoversAndCompletes(
+      self, num_elements, num_workers, num_repetitions):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = ds.repeat(num_repetitions)
+    self.evaluate(distributed_save_op.distributed_save(
+        ds, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset,
+        list(range(num_elements)) * num_repetitions,
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 5], num_sources=[1, 3])))
+  def testNonrepeatedDatasetDoesntProduceSecondRepetitionDir(
+      self, num_workers, num_sources):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = self._get_dataset(dataset_range=1000, num_sources=num_sources)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+    for stream_idx in range(num_workers):
+      for source_idx in range(num_sources):
+        self.assertFalse(
+            os.path.exists(
+                os.path.join(
+                    snapshot_dir.full_path,
+                    "streams",
+                    f"stream_{stream_idx}",
+                    "splits",
+                    f"source_{source_idx}",
+                    "repetition_1")))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def testMultipleDatasetRecoversAndCompletes(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset1 = dataset_ops.Dataset.range(1000)
+    datasets = [
+        dataset_ops.Dataset.from_tensors("a").repeat(50),
+        dataset_ops.Dataset.from_tensors("b").repeat(50),
+        dataset_ops.Dataset.from_tensors("c").repeat(50)]
+    choice_dataset = dataset_ops.Dataset.range(3).repeat()
+    dataset2 = dataset_ops.Dataset.choose_from_datasets(
+        datasets, choice_dataset)
+
+    snapshot_path1 = os.path.join(snapshot_dir.full_path, "snapshot1")
+    snapshot_path2 = os.path.join(snapshot_dir.full_path, "snapshot2")
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset1, snapshot_path1, cluster.dispatcher_address()))
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset2, snapshot_path2, cluster.dispatcher_address()))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_path1)
+    wait_for_snapshot(snapshot_path2)
+
+    dataset = dataset_ops.Dataset.load(snapshot_path1)
+    self.assertDatasetProduces(
+        dataset, list(range(1000)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def testNestedDataset(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.from_tensor_slices(range(100))
+    def interleave_fn(x):
+      ds = dataset_ops.Dataset.from_tensor_slices(range(x))
+      def flat_map_fn(y):
+        return dataset_ops.Dataset.from_tensor_slices([y])
+      return ds.flat_map(flat_map_fn)
+    dataset = dataset.interleave(
+        interleave_fn, cycle_length=2, num_parallel_calls=2)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    for i in range(num_workers):
+      cluster.stop_worker(i)
+      cluster.restart_dispatcher()
+      cluster.restart_worker(i)
+    wait_for_snapshot(snapshot_dir.full_path)
+    self.assertTrue(snapshot_is_done(snapshot_dir.full_path))
+
+  def _get_dataset(self, dataset_range=10, num_sources=1):
+    dataset = dataset_ops.Dataset.range(dataset_range)
+    if num_sources > 1:
+      dataset = dataset_ops.Dataset.zip((dataset,) * num_sources)
+    return dataset
+
+  def _splits_dir(self, snapshot_path, stream_idx=0, worker=0):
+    stream_name = f"stream_{stream_idx}"
+    self._make_stream_dir(snapshot_path, stream_name, worker=worker)
+    return os.path.join(snapshot_path, "streams", stream_name, "splits")
+
+  def _source_dir(self, snapshot_path, stream_idx=0, source_idx=0, worker=0):
+    return os.path.join(
+        self._splits_dir(snapshot_path, stream_idx, worker=worker),
+        f"source_{source_idx}",
+        "repetition_0")
+
+  def _make_stream_dir(self, snapshot_path, stream_name, worker=0):
+    stream_dir = os.path.join(snapshot_path, "streams", stream_name)
+    os.makedirs(stream_dir)
+    pathlib.Path(os.path.join(stream_dir, "owner_worker")).write_text(
+        f"{worker}")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/distributed_save_load_ft_test.py

@@ -0,0 +1,267 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Fault tolerance tests for distributed save/load."""
+
+import multiprocessing
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distributed_save_op
+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 DistributedSaveLoadFtTest(
+    data_service_test_base.TestBase, parameterized.TestCase
+):
+  """Fault tolerance tests for distributed save/load."""
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(
+              num_elements=[200],
+              num_workers=[1, 2],
+              save_repetitions=[1, 2],
+              load_repetitions=[1, 2],
+              sharding_policy=[
+                  data_service_ops.ShardingPolicy.OFF,
+                  data_service_ops.ShardingPolicy.DYNAMIC,
+              ],
+          ),
+      )
+  )
+  def test_dispatcher_restart(
+      self,
+      num_workers: int,
+      num_elements: int,
+      save_repetitions: int,
+      load_repetitions: int,
+      sharding_policy: data_service_ops.ShardingPolicy,
+  ):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    if save_repetitions > 1:
+      dataset = dataset.repeat(save_repetitions)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()
+        )
+    )
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    if load_repetitions > 1:
+      dataset = dataset.repeat(load_repetitions)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            sharding_policy,
+            cluster.dispatcher_address(),
+            max_outstanding_requests=1,
+        )
+    )
+
+    iterator = self.getNext(dataset)
+    output = [self.evaluate(iterator())]
+    cluster.restart_dispatcher()
+    output.extend(self.getIteratorOutput(iterator))
+
+    # For no sharding, dispatcher restarts do not affect data processing
+    # happening at the workers.
+    repetitions = save_repetitions * load_repetitions
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      expected = list(range(num_elements)) * repetitions * num_workers
+      self.assertCountEqual(output, expected)
+
+    # Dynamic sharding may lose splits if the dispatcher fails.
+    if sharding_policy == data_service_ops.ShardingPolicy.DYNAMIC:
+      self.assertNotEmpty(output)
+      self.assertContainsSubset(output, range(num_elements))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(
+              num_elements=[200],
+              num_workers=[1, 2],
+              save_repetitions=[1, 2],
+              load_repetitions=[1, 2],
+              sharding_policy=[
+                  # Enable dynamic sharding. Need to fix the race condition
+                  # where workers restart before sending the final task
+                  # completion update.
+                  data_service_ops.ShardingPolicy.OFF
+              ],
+          ),
+      )
+  )
+  def test_dispatcher_and_worker_restart(
+      self,
+      num_elements: int,
+      num_workers: int,
+      save_repetitions: int,
+      load_repetitions: int,
+      sharding_policy: data_service_ops.ShardingPolicy,
+  ):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    if save_repetitions > 1:
+      dataset = dataset.repeat(save_repetitions)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()
+        )
+    )
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    if load_repetitions > 1:
+      dataset = dataset.repeat(load_repetitions)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            sharding_policy,
+            cluster.dispatcher_address(),
+            max_outstanding_requests=1,
+        )
+    )
+
+    iterator = self.getNext(dataset)
+    output = [self.evaluate(iterator())]
+    for i in range(num_workers):
+      cluster.restart_dispatcher()
+      cluster.workers[i].restart()
+    output.extend(self.getIteratorOutput(iterator))
+
+    # If the sharding policy is OFF, the restarted worker will produce elements
+    # from the beginning of the dataset. The result is a partial range plus
+    # `num_elements` repetitions.
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      repetitions = save_repetitions * load_repetitions
+      self.assertContainsSubsequence(
+          sorted(output),
+          sorted(list(range(num_elements)) * repetitions * num_workers),
+      )
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(
+              load_repetitions=[1, 2],
+              sharding_policy=[
+                  data_service_ops.ShardingPolicy.OFF,
+                  data_service_ops.ShardingPolicy.DYNAMIC,
+              ],
+          ),
+      )
+  )
+  def test_add_worker_midjob(
+      self,
+      load_repetitions: int,
+      sharding_policy: data_service_ops.ShardingPolicy,
+  ):
+    num_elements = 2 * multiprocessing.cpu_count() + 100
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()
+        )
+    )
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    dataset = dataset.repeat(load_repetitions)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            sharding_policy,
+            cluster.dispatcher_address(),
+            max_outstanding_requests=1,
+        )
+    )
+    expected = list(range(num_elements)) * load_repetitions
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      expected *= 2
+
+    # Reads halfway through the dataset.
+    iterator = self.getNext(dataset)
+    output = [self.evaluate(iterator()) for _ in range(num_elements // 2)]
+
+    # Waits for a new worker to register with the dispatcher.
+    cluster.add_worker()
+    while cluster.num_registered_workers() < 2:
+      time.sleep(10 / 1000)  # 10ms
+
+    output.extend(self.getIteratorOutput(iterator))
+    self.assertCountEqual(output, expected)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(
+              num_workers=[1, 2],
+              num_elements=[200],
+              load_repetitions=[1, 2],
+              sharding_policy=[
+                  data_service_ops.ShardingPolicy.OFF,
+                  data_service_ops.ShardingPolicy.DYNAMIC,
+              ],
+          ),
+      )
+  )
+  def test_new_dataset_after_restart(
+      self,
+      num_workers: int,
+      num_elements: int,
+      load_repetitions: int,
+      sharding_policy: data_service_ops.ShardingPolicy,
+  ):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()
+        )
+    )
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    dataset = dataset.repeat(load_repetitions)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            sharding_policy, cluster.dispatcher_address()
+        )
+    )
+
+    expected = list(range(num_elements)) * load_repetitions
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      expected *= num_workers
+
+    # Re-processes the dataset after dispatcher/worker restarts.
+    cluster.restart_dispatcher()
+    cluster.workers[0].restart()
+    self.assertCountEqual(self.getDatasetOutput(dataset), expected)
+
+    cluster.restart_dispatcher()
+    cluster.workers[0].restart()
+    self.assertCountEqual(self.getDatasetOutput(dataset), expected)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/distributed_save_load_test.py

@@ -0,0 +1,497 @@
+# Copyright 2023 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 distributed save/load with the new load algorithm."""
+
+import multiprocessing
+import os
+import threading
+import time
+from typing import Callable, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distributed_save_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 load_op
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class DistributedSaveLoadTest(
+    data_service_test_base.TestBase, parameterized.TestCase):
+  """Tests for distributed save/load with the new load algorithm."""
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3],
+              num_elements=[0, 10],
+              num_repetitions=[1, 3],
+              compression=[None, "AUTO", "GZIP"],
+              max_chunk_size_bytes=[1, 16 << 10])))
+  def test_save_load(
+      self,
+      num_workers: int,
+      num_elements: int,
+      num_repetitions: int,
+      compression: Optional[str],
+      max_chunk_size_bytes: int):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers,
+        snapshot_max_chunk_size_bytes=max_chunk_size_bytes)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.repeat(num_repetitions)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset,
+        list(range(num_elements)) * num_repetitions,
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def test_concurrent_save_load(self, num_workers: int):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+
+    def load_thread_fn():
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+      self.assertDatasetProduces(
+          dataset, list(range(10)), 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(10)
+      self.evaluate(
+          distributed_save_op.distributed_save(
+              dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    save_thread = threading.Thread(target=save_thread_fn, name="save_thread")
+    save_thread.start()
+    save_thread.join()
+    load_thread.join()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 5],
+              num_elements=[10],
+              num_repetitions=[10],
+              max_chunk_size_bytes=[1, 16 << 10])))
+  def test_deterministic_load_order(
+      self,
+      num_workers: int,
+      num_elements: int,
+      num_repetitions: int,
+      max_chunk_size_bytes: int):
+    """Verifies `load` produces data deterministically after `save` finishes."""
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers,
+        snapshot_max_chunk_size_bytes=max_chunk_size_bytes)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements).shuffle(
+        buffer_size=num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    dataset = dataset.repeat(num_repetitions)
+    output = self.getDatasetOutput(dataset)
+    output_per_repetition = [
+        output[i : i + num_elements]
+        for i in range(0, len(output), num_elements)]
+    self.assertLen(output_per_repetition, num_repetitions)
+    for i in range(2, num_repetitions):  # Starts from the second repetition.
+      self.assertEqual(output_per_repetition[i], output_per_repetition[i - 1])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_elements=[20],
+              num_repetitions=[10])))
+  def test_shuffle_chunks(
+      self, num_elements: int, num_repetitions: int):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=5, snapshot_max_chunk_size_bytes=8)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    def _interleave_shuffled_chunks(
+        datasets: dataset_ops.Dataset) -> dataset_ops.Dataset:
+      # Does not shuffle when the snapshot is being written (first repetition).
+      # Otherwise, loading will be blocked to fill the shuffle buffer, not to
+      # read the partial chunks.
+      datasets = cond.cond(
+          math_ops.greater(datasets.cardinality(), 0),
+          lambda: datasets.shuffle(buffer_size=datasets.cardinality()),
+          lambda: datasets)
+      return datasets.interleave(
+          lambda x: x,
+          cycle_length=multiprocessing.cpu_count(),
+          num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    dataset = dataset_ops.Dataset.range(num_repetitions).flat_map(
+        lambda _: dataset_ops.Dataset.load(
+            snapshot_dir.full_path,
+            reader_func=_interleave_shuffled_chunks,
+            wait=True))
+
+    output = self.getDatasetOutput(dataset)
+    output_per_repetition = [
+        output[i : i + num_elements]
+        for i in range(0, len(output), num_elements)]
+    self.assertLen(output_per_repetition, num_repetitions)
+    for i in range(1, num_repetitions):
+      self.assertNotEqual(output_per_repetition[i],
+                          output_per_repetition[i - 1])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3],
+              num_elements=[0, 10],
+              repeated_load=[1, 5],
+              sharding_policy=[
+                  data_service_ops.ShardingPolicy.OFF,
+                  data_service_ops.ShardingPolicy.DYNAMIC])))
+  def test_distributed_load(
+      self,
+      num_workers: int,
+      num_elements: int,
+      repeated_load: int,
+      sharding_policy: data_service_ops.ShardingPolicy):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    if repeated_load > 1:
+      dataset = dataset.repeat(repeated_load)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            sharding_policy, cluster.dispatcher_address()))
+    expected = list(range(num_elements)) * repeated_load
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      expected *= num_workers
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def test_save_before_sample(self, num_workers: int):
+    num_elements = 10
+    num_datasets = 3
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    datasets = [
+        dataset_ops.Dataset.range(num_elements) for i in range(num_datasets)]
+    for i, dataset in enumerate(datasets):
+      self.evaluate(
+          distributed_save_op.distributed_save(
+              dataset,
+              os.path.join(snapshot_dir.full_path, f"dataset_{i}"),
+              cluster.dispatcher_address()))
+
+    loaded_datasets = []
+    for i in range(len(datasets)):
+      snapshot_path = os.path.join(snapshot_dir.full_path, f"dataset_{i}")
+      loaded_datasets.append(dataset_ops.Dataset.load(snapshot_path, wait=True))
+    dataset = dataset_ops.Dataset.sample_from_datasets(
+        loaded_datasets,
+        weights=[1.0] * num_datasets,
+        stop_on_empty_dataset=False)
+    self.assertDatasetProduces(
+        dataset,
+        list(range(num_elements)) * num_datasets,
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3], num_repetitions=[1, 3])))
+  def test_save_after_sample(self, num_workers: int, num_repetitions: int):
+    num_elements = 10
+    num_datasets = 3
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    datasets = [
+        dataset_ops.Dataset.range(num_elements) for i in range(num_datasets)]
+    if num_repetitions > 1:
+      datasets = [dataset.repeat(num_repetitions) for dataset in datasets]
+    dataset = dataset_ops.Dataset.sample_from_datasets(
+        datasets, weights=[1.0] * num_datasets, stop_on_empty_dataset=False)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset,
+        list(range(num_elements)) * num_datasets * num_repetitions,
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def test_enumerate(self, num_workers: int):
+    cluster = data_service_test_base.TestCluster(num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.from_tensor_slices(["a", "b", "c"])
+    dataset = dataset.repeat(3)
+    dataset = dataset.enumerate()
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    indexes, elements = map(list, zip(*self.getDatasetOutput(dataset)))
+    if num_workers == 1:
+      self.assertCountEqual(indexes, list(range(9)))
+    self.assertCountEqual(elements, [b"a", b"b", b"c"] * 3)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_worker_failure(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    components = np.array([1.0, 2.0, 3.0, np.nan, 5.0]).astype(np.float32)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.map(lambda x: array_ops.check_numerics(x, "message"))
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_elements=[10])))
+  def test_dataset_spec_file_is_optional(self, num_elements: int):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+    # After removing the dataset_spec file, the loaded dataset should produce
+    # the same output.
+    os.remove(os.path.join(
+        snapshot_dir.full_path, dataset_ops.DATASET_SPEC_FILENAME))
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_elements=[10])))
+  def test_empty_dataset_spec_file(self, num_elements: int):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+    dataset_spec_file = os.path.join(
+        snapshot_dir.full_path, dataset_ops.DATASET_SPEC_FILENAME)
+    with open(dataset_spec_file, "w") as f:
+      f.write("")
+
+    # Reads element_spec from the metadata file.
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def test_snapshot_does_not_exist(self):
+    snapshot_dir = data_service_test_base.TempDir()
+    with self.assertRaises(errors.NotFoundError):
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=False)
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_repetitions=[None, 0, 1, 3])))
+  def test_snapshot_chunks_cardinality(self, num_repetitions: int):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = load_op._ListSnapshotChunksDataset(snapshot_dir.full_path)
+    if num_repetitions != 1:
+      dataset = dataset.repeat(num_repetitions)
+    while self.evaluate(dataset.cardinality()) == dataset_ops.UNKNOWN:
+      time.sleep(.1)
+
+    num_chunks = len(os.listdir(os.path.join(snapshot_dir.full_path, "chunks")))
+    expected_cardinality = (
+        dataset_ops.INFINITE
+        if num_repetitions is None
+        else num_chunks * num_repetitions)
+    self.assertEqual(self.evaluate(dataset.cardinality()), expected_cardinality)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_repetitions=[50])))
+  def test_snapshot_chunks_order(self, num_repetitions: int):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1, snapshot_max_chunk_size_bytes=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(1)
+    dataset = dataset.repeat(num_repetitions)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = load_op._ListSnapshotChunksDataset(snapshot_dir.full_path)
+    while self.evaluate(dataset.cardinality()) == dataset_ops.UNKNOWN:
+      time.sleep(.1)
+
+    chunk_indices = [
+        int(os.path.basename(str(chunk)).split("_")[2])
+        for chunk in self.getDatasetOutput(dataset)]
+    self.assertEqual(chunk_indices, sorted(chunk_indices),
+                     "Snapshot chunks should be sorted by chunk indices.")
+
+
+class SaveLoadCheckpointTest(
+    data_service_test_base.TestBase,
+    checkpoint_test_base.CheckpointTestBase,
+    parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3],
+              num_elements=[0, 10],
+              num_repetitions=[1, 5])))
+  def test_save_load_checkpoint(
+      self,
+      verify_fn: Callable[..., None],
+      num_workers: int,
+      num_elements: int,
+      num_repetitions: int,):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    def _build_ds() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+      if num_repetitions > 1:
+        dataset = dataset.repeat(num_repetitions)
+      return dataset
+
+    # Compares output ignoring order since the first repetition may be
+    # non-deterministic.
+    verify_fn(
+        self,
+        _build_ds,
+        num_outputs=num_elements * num_repetitions,
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3],
+              num_elements=[0, 10],
+              num_repetitions=[5],
+              max_chunk_size_bytes=[1, 16 << 10])))
+  def test_skip_first_repetition(
+      self,
+      verify_fn: Callable[..., None],
+      num_workers: int,
+      num_elements: int,
+      num_repetitions: int,
+      max_chunk_size_bytes: int):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers,
+        snapshot_max_chunk_size_bytes=max_chunk_size_bytes)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    def _build_ds() -> dataset_ops.Dataset:
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path, wait=True)
+      dataset = dataset.repeat(num_repetitions)
+      # Skips the first repetition. The remaining repetitions should be
+      # deterministic.
+      dataset = dataset.skip(num_elements)
+      return dataset
+
+    verify_fn(
+        self,
+        _build_ds,
+        num_outputs=num_elements * (num_repetitions - 1),
+        assert_items_equal=False)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/distributed_save_test.py

@@ -0,0 +1,406 @@
+# Copyright 2022 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.distributed_save."""
+
+import os
+import time
+
+from absl.testing import parameterized
+
+import numpy as np
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distributed_save_op
+from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
+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 DistributedSaveTest(
+    data_service_test_base.TestBase,
+    parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3], num_elements=[0, 10, 1000])))
+  def testSaveLoad(self, num_workers, num_elements):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(num_elements)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(compression=[None, "AUTO", "GZIP"])))
+  def testCompression(self, compression):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset,
+        snapshot_dir.full_path,
+        cluster.dispatcher_address(),
+        compression=compression))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset, list(range(10)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3], num_repetitions=[1, 5])))
+  def testRepeatedDataset(self, num_workers, num_repetitions):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(1000)
+    dataset = dataset.repeat(num_repetitions)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset,
+        snapshot_dir.full_path,
+        cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset, list(range(1000)) * num_repetitions, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromDatasets(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(["a", "a", "a", "a", "a"]),
+        dataset_ops.Dataset.from_tensor_slices(["b", "b", "b", "b", "b"]),
+        dataset_ops.Dataset.from_tensor_slices(["c", "c", "c", "c", "c"])]
+    choice_dataset = dataset_ops.Dataset.range(3).repeat()
+    dataset = dataset_ops.Dataset.choose_from_datasets(datasets, choice_dataset)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset, [b"a", b"b", b"c"] * 5, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testChooseFromRepeatedDatasets(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    datasets = [
+        dataset_ops.Dataset.from_tensors("a").repeat(5),
+        dataset_ops.Dataset.from_tensors("b").repeat(5),
+        dataset_ops.Dataset.from_tensors("c").repeat(10)]
+    choice_dataset = dataset_ops.Dataset.range(3).repeat()
+    dataset = dataset_ops.Dataset.choose_from_datasets(
+        datasets, choice_dataset, stop_on_empty_dataset=False)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    self.assertDatasetProduces(
+        dataset, [b"a", b"b", b"c"] * 5 + [b"c"] * 5, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(num_workers=[1, 3])))
+  def testWriteMultipleDatasets(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset1 = dataset_ops.Dataset.range(100)
+    datasets = [
+        dataset_ops.Dataset.from_tensors("a").repeat(5),
+        dataset_ops.Dataset.from_tensors("b").repeat(5),
+        dataset_ops.Dataset.from_tensors("c").repeat(5)]
+    choice_dataset = dataset_ops.Dataset.range(3).repeat()
+    dataset2 = dataset_ops.Dataset.choose_from_datasets(
+        datasets, choice_dataset)
+
+    snapshot_path1 = os.path.join(snapshot_dir.full_path, "snapshot1")
+    snapshot_path2 = os.path.join(snapshot_dir.full_path, "snapshot2")
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset1, snapshot_path1, cluster.dispatcher_address()))
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset2, snapshot_path2, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_path1)
+    _wait_for_snapshot(snapshot_path2)
+
+    dataset1 = dataset_ops.Dataset.load(snapshot_path1)
+    self.assertDatasetProduces(
+        dataset1, list(range(100)), assert_items_equal=True)
+    self.assertDatasetProduces(
+        dataset2, [b"a", b"b", b"c"] * 5, assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3], snapshot_max_chunk_size_bytes=[1, 100])))
+  def testLoadWithCustomReaderFunc(
+      self, num_workers, snapshot_max_chunk_size_bytes):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=num_workers,
+        snapshot_max_chunk_size_bytes=snapshot_max_chunk_size_bytes)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    def custom_reader_func(datasets):
+      datasets = datasets.shuffle(3)
+      return datasets.interleave(
+          lambda x: x, num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    dataset = dataset_ops.Dataset.load(
+        snapshot_dir.full_path, reader_func=custom_reader_func)
+    self.assertDatasetProduces(
+        dataset, list(range(10)), assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_workers=[1, 3],
+              repeated_load=[1, 5],
+              sharding_policy=[
+                  data_service_ops.ShardingPolicy.OFF,
+                  data_service_ops.ShardingPolicy.DYNAMIC])))
+  def testDistributedLoad(self, num_workers, repeated_load, sharding_policy):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    if repeated_load > 1:
+      dataset = dataset.repeat(repeated_load)
+    dataset = dataset.apply(
+        data_service_ops.distribute(
+            processing_mode=sharding_policy,
+            service=cluster.dispatcher_address()))
+
+    expected = list(range(10)) * repeated_load
+    if sharding_policy == data_service_ops.ShardingPolicy.OFF:
+      expected *= num_workers
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImbalancedZipAndRepeat(self):
+    smaller_num_elements = 200
+    larger_num_elements = 1000
+    repetitions = 3
+
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset1 = dataset_ops.Dataset.range(smaller_num_elements)
+    dataset2 = dataset_ops.Dataset.range(larger_num_elements)
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    dataset = dataset.repeat(repetitions)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+    expected = repetitions * (
+        list(zip(range(smaller_num_elements), range(smaller_num_elements))))
+    self.assertDatasetProduces(dataset, expected, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSnapshotDoesNotExist(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    with self.assertRaises(errors.NotFoundError):
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+      dataset = dataset.apply(
+          data_service_ops.distribute(
+              data_service_ops.ShardingPolicy.OFF,
+              cluster.dispatcher_address()))
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDuplicateSnapshot(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        errors.AlreadyExistsError, "already started or completed"):
+      self.evaluate(
+          distributed_save_op.distributed_save(
+              dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+      self.evaluate(
+          distributed_save_op.distributed_save(
+              dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWorkerFailure(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    components = np.array([1.0, 2.0, 3.0, np.nan, 5.0]).astype(np.float32)
+    dataset = dataset_ops.Dataset.from_tensor_slices(components)
+    dataset = dataset.map(lambda x: array_ops.check_numerics(x, "message"))
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_error(snapshot_dir.full_path)
+
+    with self.assertRaisesRegex(
+        ValueError, "The save job failed to write it."):
+      dataset = dataset_ops.Dataset.load(snapshot_dir.full_path)
+      self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidSnapshotPath(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Snapshot path cannot contain path traversal characters"):
+      self.evaluate(distributed_save_op.distributed_save(
+          dataset,
+          "../invalid_dir",
+          cluster.dispatcher_address()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBadDispatcherAddress(self):
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegex(ValueError, "must be a string"):
+      self.evaluate(distributed_save_op.distributed_save(dataset, "", 1))
+    with self.assertRaisesRegex(ValueError, "must not be empty"):
+      self.evaluate(distributed_save_op.distributed_save(dataset, "", ""))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBadCardinality(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10).repeat()
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Saving an infinite dataset is not allowed"):
+      self.evaluate(distributed_save_op.distributed_save(
+          dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBadElementSpec(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path,
+        cluster.dispatcher_address(),
+        compression="AUTO"))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "User specified element_spec bad_element_spec, but the actual "
+        "element_spec is TensorSpec"):
+      _ = dataset_ops.Dataset.load(snapshot_dir.full_path,
+                                   element_spec="bad_element_spec")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBadCompression(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(distributed_save_op.distributed_save(
+        dataset, snapshot_dir.full_path,
+        cluster.dispatcher_address(),
+        compression="AUTO"))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "User specified compression ZLIB, but the actual compression is "
+        "SNAPPY."):
+      _ = dataset_ops.Dataset.load(snapshot_dir.full_path, compression="ZLIB")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRequiresFaultTolerantMode(self):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1, fault_tolerant_mode=False)
+    snapshot_dir = data_service_test_base.TempDir()
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "tf.data distributed snapshot requires running tf.data service in the "
+        "fault tolerant mode."):
+      self.evaluate(distributed_save_op.distributed_save(
+          dataset_ops.Dataset.range(10), snapshot_dir.full_path,
+          cluster.dispatcher_address(),
+          compression="AUTO"))
+
+
+class LoadCheckpointTest(
+    data_service_test_base.TestBase,
+    checkpoint_test_base.CheckpointTestBase,
+    parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          checkpoint_test_base.default_test_combinations(),
+      )
+  )
+  def testLoadCheckpoint(self, verify_fn):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    snapshot_dir = data_service_test_base.TempDir()
+    dataset = dataset_ops.Dataset.range(10)
+    self.evaluate(
+        distributed_save_op.distributed_save(
+            dataset, snapshot_dir.full_path, cluster.dispatcher_address()))
+    _wait_for_snapshot(snapshot_dir.full_path)
+
+    def _build_ds():
+      return dataset_ops.Dataset.load(snapshot_dir.full_path)
+
+    verify_fn(self, _build_ds, num_outputs=10, assert_items_equal=True)
+
+
+def _wait_for_snapshot(snapshot_path):
+  while not os.path.exists(
+      _pywrap_snapshot_utils.TF_DATA_SnapshotDoneFilePath(snapshot_path)):
+    time.sleep(0.1)
+
+
+def _wait_for_error(snapshot_path):
+  while not os.path.exists(
+      _pywrap_snapshot_utils.TF_DATA_SnapshotErrorFilePath(snapshot_path)):
+    time.sleep(0.1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/dynamic_sharding_test.py

@@ -0,0 +1,649 @@
+# 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 dynamic sharding."""
+import collections
+import time
+
+from absl import logging
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.protobuf import service_config_pb2
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.service import server_lib
+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 readers
+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.platform import test
+
+
+class DynamicShardingTest(data_service_test_base.TestBase,
+                          parameterized.TestCase):
+
+  def _make_dynamic_sharding_dataset(self, dataset, cluster):
+    return self.make_distributed_dataset(
+        dataset,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+        job_name="job_name")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testBasic(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(
+        ds, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoJobName(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self.make_distributed_dataset(
+        ds,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+        job_name=None)
+    self.assertDatasetProduces(
+        ds, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testTensorSlices(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    vals = [5, 1, 2, 4]
+    ds = dataset_ops.Dataset.from_tensor_slices(vals)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(ds, vals, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInterleave(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    elements = [1, 5, 0]
+    ds = dataset_ops.Dataset.from_tensor_slices(elements)
+    ds = ds.interleave(lambda x: dataset_ops.Dataset.from_tensor_slices([x]))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(ds, elements, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testParallelInterleave(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    elements = [1, 5, 0]
+    ds = dataset_ops.Dataset.from_tensor_slices(elements)
+    ds = ds.interleave(
+        lambda x: dataset_ops.Dataset.from_tensor_slices([x]),
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(ds, elements, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlatMap(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    elements = [1, 5, 0]
+    ds = dataset_ops.Dataset.from_tensor_slices(elements)
+    ds = ds.flat_map(lambda x: dataset_ops.Dataset.from_tensor_slices([x]))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(ds, elements, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testGroupByWindow(self):
+    # Verify that split providers are not propagated into iterators created for
+    # the reduce datasets created by the reduce_fn in group_by_window.
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    elements = [1, 5, 0]
+    ds = dataset_ops.Dataset.from_tensor_slices(elements)
+
+    def reduce_fn(_, window):
+      return dataset_ops.Dataset.zip((window, dataset_ops.Dataset.range(100)))
+
+    ds = ds.group_by_window(lambda x: 0, reduce_fn, window_size=3)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    # This will fail if the tensor_slices split provider is propagated into the
+    # `reduce_fn`, since the `zip` requires either 0 or 2 split providers.
+    self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRepeatBeforeDistribution(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_repeats = 5
+    num_elements = 20
+    ds = dataset_ops.Dataset.range(num_elements).repeat(num_repeats)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(
+        ds, num_repeats * list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRepeatAfterDistribution(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_repeats = 5
+    num_elements = 20
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    ds = ds.repeat(num_repeats)
+    self.assertDatasetProduces(
+        ds, num_repeats * list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testForeverRepeat(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_elements = 20
+    elements_to_read = 1000
+    ds = dataset_ops.Dataset.range(num_elements).repeat()
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    get_next = self.getNext(ds)
+    results = {}
+    for _ in range(elements_to_read):
+      val = self.evaluate(get_next())
+      if val not in results:
+        results[val] = 0
+      results[val] += 1
+    for i in range(num_elements):
+      self.assertGreater(results[i], elements_to_read / num_elements / 2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testForeverRepeatFewElements(self):
+    num_workers = 5
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    # Less than the number of workers, so that some workers get zero elements on
+    # the first repetition.
+    num_elements = 1
+    ds = dataset_ops.Dataset.range(num_elements).repeat()
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    get_next = self.getNext(ds)
+    for _ in range(20):
+      self.assertEqual(self.evaluate(get_next()), 0)
+
+    # Stop all but one worker and check that we can still read.
+    for i in range(num_workers - 1):
+      cluster.workers[i].stop()
+    for _ in range(20):
+      self.assertEqual(self.evaluate(get_next()), 0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleAndRepeat(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_repeats = 5
+    num_elements = 20
+    ds = dataset_ops.Dataset.range(num_elements).shuffle(num_elements).repeat(
+        num_repeats)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    self.assertDatasetProduces(
+        ds, num_repeats * list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZip(self):
+    num_elements = 10
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    a = dataset_ops.Dataset.range(num_elements)
+
+    ds = dataset_ops.Dataset.zip((a, a))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    self.assertDatasetProduces(
+        ds, list(zip(range(num_elements), range(num_elements))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNestedZip(self):
+    num_elements = 10
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    a = dataset_ops.Dataset.range(num_elements)
+
+    ds = dataset_ops.Dataset.zip((a, a))
+    ds = dataset_ops.Dataset.zip((a, a, ds, a))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    b = list(range(10))
+    self.assertDatasetProduces(ds, list(zip(b, b, zip(b, b), b)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImbalancedZip(self):
+    smaller_num_elements = 200
+    larger_num_elements = 1000
+
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    a = dataset_ops.Dataset.range(smaller_num_elements)
+    b = dataset_ops.Dataset.range(larger_num_elements)
+
+    ds = dataset_ops.Dataset.zip((a, b))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    self.assertDatasetProduces(
+        ds, list(zip(range(smaller_num_elements), range(smaller_num_elements))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImbalancedZipAndRepeat(self):
+    smaller_num_elements = 200
+    larger_num_elements = 1000
+    repetitions = 3
+
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    a = dataset_ops.Dataset.range(smaller_num_elements)
+    b = dataset_ops.Dataset.range(larger_num_elements)
+
+    ds = dataset_ops.Dataset.zip((a, b))
+    ds = ds.repeat(repetitions)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    expected = repetitions * (
+        list(zip(range(smaller_num_elements), range(smaller_num_elements)))
+    )
+
+    self.assertDatasetProduces(ds, expected)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImbalancedZipMultiWorker(self):
+    smaller_num_elements = 200
+    larger_num_elements = 1000
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    a = dataset_ops.Dataset.range(smaller_num_elements)
+    b = dataset_ops.Dataset.range(larger_num_elements)
+
+    ds = dataset_ops.Dataset.zip((a, b))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    # Cannot assert specific elements because the range datasets are split
+    # nondeterministically and may not line up.
+    self.assertLen(self.getDatasetOutput(ds), smaller_num_elements)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipDifferentRates(self):
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    a = dataset_ops.Dataset.range(100)
+    b = dataset_ops.Dataset.range(100).filter(
+        lambda x: math_ops.equal(x % 10, 0))
+
+    ds = dataset_ops.Dataset.zip((a, b))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    self.assertLen(self.getDatasetOutput(ds), 10)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipDifferentRepeats(self):
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    a = dataset_ops.Dataset.range(50)
+    b = dataset_ops.Dataset.range(10).repeat(10)
+
+    ds = dataset_ops.Dataset.zip((a, b))
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    self.assertLen(self.getDatasetOutput(ds), 50)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSampleFromDatasets(self):
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    num_samples = 200
+    weights = [.6, .3, .1]
+    classes = len(weights)
+
+    # Create a dataset that samples each integer in `[0, num_datasets)`
+    # with probability given by `weights[i]`.
+    ds = dataset_ops.Dataset.sample_from_datasets(
+        [dataset_ops.Dataset.from_tensors(i).repeat() for i in range(classes)],
+        weights)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    ds = ds.take(num_samples)
+
+    freqs = np.zeros([classes])
+    for v in self.getDatasetOutput(ds, requires_initialization=True):
+      freqs[v] += 1
+
+    self.assertGreater(freqs[0], freqs[1])
+    self.assertGreater(freqs[1], freqs[2])
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_workers=[1, 3])))
+  def testChooseFromDatasets(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    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)
+    ds = dataset_ops.Dataset.choose_from_datasets(datasets, choice_dataset)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    expected = [words[i] for i in choice_array] * num_workers
+
+    assert_items_equal = (num_workers > 1)
+    self.assertDatasetProduces(
+        ds, expected, assert_items_equal=assert_items_equal)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEnumerateReplicateOnSplit(self):
+    num_workers = 3
+    cluster = data_service_test_base.TestCluster(num_workers)
+    ds = dataset_ops.Dataset.from_tensor_slices(["a", "b", "c"]).repeat()
+    ds = ds.enumerate()
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    get_next = self.getNext(ds)
+
+    counts = collections.defaultdict(int)
+    while True:
+      i, _ = self.evaluate(get_next())
+      counts[i] += 1
+      # Read until all workers have reached enumeration index 10.
+      if counts[10] == num_workers:
+        break
+
+    for i in range(10):
+      self.assertEqual(counts[i], num_workers)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZipReplicateOnSplit(self):
+    num_workers = 3
+    cluster = data_service_test_base.TestCluster(num_workers)
+
+    ds1 = dataset_ops.Dataset.range(100, 1000, output_type=dtypes.int32)
+    ds2 = dataset_ops.Dataset.from_tensor_slices(range(0, 5))
+    ds2 = dataset_ops.apply_rewrite(ds2, "replicate_on_split")
+
+    ds = dataset_ops.Dataset.zip(ds1, ds2)
+    ds = ds.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="shared_job",
+        )
+    )
+    ds2_output = [y for x, y in self.getDatasetOutput(ds)]
+    self.assertCountEqual(ds2_output, [0, 1, 2, 3, 4] * num_workers)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_workers=[1, 3])))
+  def testConcatenate(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    a = dataset_ops.Dataset.range(100)
+    b = dataset_ops.Dataset.range(100, 200)
+    ds = a.concatenate(b)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    assert_items_equal = (num_workers > 1)
+    self.assertDatasetProduces(
+        ds, list(range(200)), assert_items_equal=assert_items_equal)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_workers=[1, 3])))
+  def testTake(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    a = dataset_ops.Dataset.range(3).repeat().take(5)
+    b = dataset_ops.Dataset.range(100, 103).repeat().take(5)
+    ds = a.concatenate(b)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    assert_items_equal = num_workers > 1
+    if num_workers == 1:
+      expected = [0, 1, 2, 0, 1, 100, 101, 102, 100, 101]
+    else:
+      expected = [0, 1, 2] * 5 + [100, 101, 102] * 5
+    self.assertDatasetProduces(
+        ds, expected, assert_items_equal=assert_items_equal
+    )
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_workers=[1, 3])))
+  def testTakeWithRestartedWorker(self, num_workers):
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    a = dataset_ops.Dataset.range(3).repeat()
+    b = dataset_ops.Dataset.range(0)
+    c = dataset_ops.Dataset.range(100, 103).repeat()
+    a = a.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="job_a",
+        )
+    )
+    b = b.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="job_b",
+        )
+    )
+    c = c.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="job_c",
+        )
+    )
+    a = a.take(5)
+    b = b.take(5)
+    c = c.take(5)
+    ds = a.concatenate(b).concatenate(c)
+    ds = ds.prefetch(buffer_size=dataset_ops.AUTOTUNE)
+
+    output = []
+    get_next = self.getNext(ds)
+    for _ in range(5):
+      output.append(self.evaluate(get_next()))
+    cluster.workers[0].restart(use_same_port=True)
+
+    output.extend(self.getIteratorOutput(get_next))
+    logging.info("Dataset output: %s", output)
+
+    # 5 elements from each of a and b.
+    self.assertTrue(all(x < 100 for x in output[:5]))
+    self.assertTrue(all(x >= 100 for x in output[5:]))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(already_written=[True, False])))
+  def testSnapshot(self, already_written):
+    num_workers = 3
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    ds = dataset_ops.Dataset.range(100)
+    ds = ds.snapshot(self.get_temp_dir())
+    if already_written:
+      # Materialize the snapshot.
+      self.getDatasetOutput(ds)
+
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+    error_regex = "Splitting is not implemented for snapshot datasets"
+    with self.assertRaisesRegex(errors.UnimplementedError, error_regex):
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDistributedDataset(self):
+    cluster_1 = data_service_test_base.TestCluster(num_workers=1)
+    cluster_2 = data_service_test_base.TestCluster(num_workers=1)
+    num_sizes = 10
+    size_repeats = 5
+    numbers = [1 * i for i in range(num_sizes)] * size_repeats
+    ds = dataset_ops.Dataset.from_tensor_slices(numbers)
+    ds = self.make_distributed_dataset(
+        ds, cluster_1, processing_mode=data_service_ops.ShardingPolicy.OFF)
+    ds = ds.map(lambda x: x + 1)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster_2)
+
+    error_regex = ("Cannot create split providers for dataset " +
+                   "of type DataServiceDataset")
+    with self.assertRaisesRegex(errors.UnimplementedError, error_regex):
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDistributedEpoch(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = self.make_distributed_dataset(
+        ds, cluster, processing_mode="distributed_epoch")
+    self.assertDatasetProduces(
+        ds, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlatMapWithRepeat(self):
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    ds = dataset_ops.Dataset.range(5)
+
+    def flat_map_fn(_):
+      return dataset_ops.Dataset.from_tensor_slices(["a", "b", "c"]).repeat(10)
+
+    ds = ds.flat_map(flat_map_fn)
+    ds = self._make_dynamic_sharding_dataset(ds, cluster)
+
+    self.assertDatasetProduces(
+        ds, [b"a", b"b", b"c"] * 50, assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleRegistration(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    dataset = dataset_ops.Dataset.range(100)
+    dataset1 = dataset.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="shared_job"
+        )
+    )
+    dataset2 = dataset.apply(
+        data_service_ops.distribute(
+            data_service_ops.ShardingPolicy.DYNAMIC,
+            cluster.dispatcher_address(),
+            job_name="shared_job"
+        )
+    )
+
+    output = []
+    iter1 = self.getNext(dataset1)
+    iter2 = self.getNext(dataset2)
+    output.extend(self.getIteratorOutput(iter1))
+    output.extend(self.getIteratorOutput(iter2))
+    self.assertCountEqual(output, list(range(100)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentDatasetIdsForSameJob(self):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    dataset = dataset_ops.Dataset.range(100)
+    dataset_id1 = data_service_ops.register_dataset(
+        cluster.dispatcher_address(), dataset, dataset_id="dataset_id1"
+    )
+    dataset1 = data_service_ops.from_dataset_id(
+        data_service_ops.ShardingPolicy.DYNAMIC,
+        cluster.dispatcher_address(),
+        dataset_id1,
+        element_spec=dataset.element_spec,
+        job_name="shared_job",
+    )
+
+    dataset_id2 = data_service_ops.register_dataset(
+        cluster.dispatcher_address(), dataset, dataset_id="dataset_id2"
+    )
+    dataset2 = data_service_ops.from_dataset_id(
+        data_service_ops.ShardingPolicy.DYNAMIC,
+        cluster.dispatcher_address(),
+        dataset_id2,
+        element_spec=dataset.element_spec,
+        job_name="shared_job",
+    )
+
+    output = []
+    iter1 = self.getNext(dataset1)
+    iter2 = self.getNext(dataset2)
+    for _ in range(5):
+      output.append(self.evaluate(iter1()))
+      output.append(self.evaluate(iter2()))
+    output.extend(self.getIteratorOutput(iter1))
+    output.extend(self.getIteratorOutput(iter2))
+    # Produces 2 copies of the dataset because `from_dataset_id` prepends the
+    # dataset ID to the job name.
+    self.assertCountEqual(output, list(range(100)) * 2)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testGcDynamicShardingJobIfRequested(self):
+    dispatcher = server_lib.DispatchServer(
+        service_config_pb2.DispatcherConfig(
+            protocol="grpc",
+            job_gc_check_interval_ms=50,
+            job_gc_timeout_ms=20,
+            gc_dynamic_sharding_jobs=True,
+        )
+    )
+    dispatcher_address = dispatcher.target.split("://")[1]
+    worker = server_lib.WorkerServer(
+        server_lib.WorkerConfig(
+            dispatcher_address=dispatcher_address, heartbeat_interval_ms=100
+        )
+    )
+
+    num_elements = 1000
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.apply(
+        data_service_ops._distribute(
+            processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+            service=dispatcher.target,
+        )
+    )
+    it = iter(dataset)
+    self.assertEqual(self.evaluate(next(it)), 0)
+    self.assertEqual(worker._num_tasks(), 1)
+    del it
+    while worker._num_tasks() > 0:
+      time.sleep(0.1)
+
+    # If re-creating the iterator with the same job, the elements are revisited.
+    it = iter(dataset)
+    self.assertEqual(self.evaluate(next(it)), 0)
+    del it
+    while worker._num_tasks() > 0:
+      time.sleep(0.1)
+    del worker
+
+
+class DynamicShardingFilesTest(data_service_test_base.TestBase,
+                               tf_record_test_base.TFRecordTestBase,
+                               parameterized.TestCase):
+
+  def setUp(self):
+    super(DynamicShardingFilesTest, self).setUp()
+    self._num_files = 5
+    self._num_records = 5
+    self._filenames = self._createFiles()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffleFiles(self):
+    cluster = data_service_test_base.TestCluster(num_workers=3)
+    shuffled_filenames = random_ops.random_shuffle(self._filenames)
+    dataset = dataset_ops.Dataset.from_tensor_slices(shuffled_filenames)
+    dataset = dataset.interleave(readers.TFRecordDataset)
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC)
+    # pylint:disable=g-complex-comprehension
+    expected = [
+        b"Record %d of file %d" % (record, file)
+        for file in range(0, 5)
+        for record in range(0, 5)
+    ]
+    self.assertDatasetProduces(
+        dataset,
+        expected,
+        requires_initialization=True,
+        assert_items_equal=True)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/fault_tolerance_test.py

@@ -0,0 +1,328 @@
+# 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 service ops where servers are started late or preempted."""
+import multiprocessing
+import threading
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+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.ops import array_ops
+from tensorflow.python.platform import test
+
+TMP_WORK_DIR = data_service_test_base.TMP_WORK_DIR
+NO_WORK_DIR = data_service_test_base.NO_WORK_DIR
+
+
+class FaultToleranceTest(data_service_test_base.TestBase,
+                         parameterized.TestCase):
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherStop(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    iterator = iter(ds)
+    results = []
+    results.append(next(iterator).numpy())
+    cluster.stop_dispatcher()
+    # After the dispatcher dies, the worker should continue providing the rest
+    # of the dataset's elements.
+    for _ in range(num_elements - 1):
+      results.append(next(iterator).numpy())
+    self.assertEqual(results, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartBeforeReading(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    cluster.restart_dispatcher()
+
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartDuringReading(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    iterator = iter(ds)
+    results = []
+    for _ in range(num_elements // 2):
+      results.append(next(iterator).numpy())
+    cluster.restart_dispatcher()
+    for elem in iterator:
+      results.append(elem.numpy())
+
+    self.assertEqual(list(range(num_elements)), results)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartDuringDistributedEpoch(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC)
+    iterator = iter(ds)
+    results = []
+    for _ in range(num_elements // 2):
+      results.append(next(iterator).numpy())
+    cluster.restart_dispatcher()
+    for elem in iterator:
+      results.append(elem.numpy())
+
+    self.assertEqual(list(range(num_elements)), results)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartDuringDistributedEpochRepeat(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    repetitions = 5
+    breakpoints = [50, 250, 450, 500]
+    ds = dataset_ops.Dataset.range(num_elements)
+    ds = ds.repeat(repetitions)
+    ds = self.make_distributed_dataset(
+        ds, cluster, processing_mode=data_service_ops.ShardingPolicy.DYNAMIC)
+
+    iterator = iter(ds)
+    results = []
+    for breakpoint_ in breakpoints:
+      for _ in range(len(results), breakpoint_):
+        results.append(next(iterator).numpy())
+      cluster.restart_dispatcher()
+
+    self.assertCountEqual(repetitions * list(range(num_elements)), results)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartBetweenIterations(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(100, cluster)
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+    cluster.restart_dispatcher()
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherRestartWithMultipleDatasets(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    datasets = []
+    for _ in range(10):
+      datasets.append(self.make_distributed_range_dataset(100, cluster))
+      cluster.restart_dispatcher()
+
+    for ds in datasets:
+      self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherManyRestarts(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements_start = 10
+    num_elements_end = 15
+    datasets = []
+    for num_elements in range(num_elements_start, num_elements_end):
+      datasets.append(
+          self.make_distributed_range_dataset(num_elements, cluster))
+      cluster.restart_dispatcher()
+    for ds, num_elements in zip(datasets,
+                                range(num_elements_start, num_elements_end)):
+      self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherAndWorkerRestart(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+
+    cluster.restart_dispatcher()
+    cluster.workers[0].restart()
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+    cluster.restart_dispatcher()
+    cluster.workers[0].restart()
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testDispatcherAndMultiWorkerRestart(self):
+    num_workers = 2
+    cluster = data_service_test_base.TestCluster(num_workers=num_workers)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    iterator = iter(ds)
+    results = []
+
+    cluster.restart_dispatcher()
+    for worker_index in range(num_workers):
+      cluster.workers[worker_index].restart()
+    for elem in iterator:
+      results.append(elem.numpy())
+    self.assertCountEqual(num_workers * list(range(num_elements)), results)
+    cluster.restart_dispatcher()
+    for worker_index in range(num_workers):
+      cluster.workers[worker_index].restart()
+    for elem in iterator:
+      results.append(elem.numpy())
+    self.assertCountEqual(num_workers * list(range(num_elements)), results)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testStartServersLate(self):
+    # Test that the data service client performs retries instead of failing when
+    # the dataset is created before the master and worker are started.
+    try:
+      import portpicker  # pylint: disable=g-import-not-at-top
+      dispatcher_port = portpicker.pick_unused_port()
+    except:
+      raise self.skipTest("Flakes in portpicker library do not represent "
+                          "TensorFlow errors.")
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1, dispatcher_port=dispatcher_port, start=False)
+
+    def start_servers():
+      time.sleep(0.5)
+      cluster.start_dispatcher()
+      cluster.start_workers()
+
+    start_servers_thread = threading.Thread(target=start_servers, daemon=True)
+    start_servers_thread.start()
+
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    results = [elem.numpy() for elem in ds]
+    self.assertEqual(list(range(num_elements)), results)
+    start_servers_thread.join()
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testAddWorkerMidJob(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 2 * multiprocessing.cpu_count() + 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    iterator = iter(ds)
+    results = []
+    # Read halfway through the dataset.
+    for _ in range(num_elements // 2):
+      results.append(next(iterator).numpy())
+
+    cluster.add_worker()
+    # Wait for the new worker to register with the dispatcher.
+    while cluster.num_registered_workers() < 2:
+      time.sleep(10 / 1000)  # 10ms
+
+    for elem in iterator:
+      results.append(elem.numpy())
+
+    self.assertCountEqual(2 * list(range(num_elements)), results)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testRemoveMoreWorkersThanMaxOutstandingRequests(self):
+    num_workers = 5
+    cluster = data_service_test_base.TestCluster(num_workers)
+    num_elements = 2**55  # Effectively infinite
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, max_outstanding_requests=1)
+    iterator = iter(ds)
+    zeros_seen = 0
+    # Read until we've read from all workers. Each worker produces a zero first.
+    while zeros_seen < num_workers:
+      if next(iterator).numpy() == 0:
+        zeros_seen += 1
+
+    for i in range(num_workers - 1):
+      cluster.stop_worker(i)
+
+    # Read additional elements to make sure that stopping 4/5 workers doesn't
+    # result in a hang.
+    for _ in range(10):
+      next(iterator).numpy()
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(use_same_port=[True, False]),
+                         data_service_test_base.all_cluster_configurations()))
+  def testRestartWorker(self, use_same_port, work_dir, fault_tolerant_mode):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1,
+        work_dir=work_dir,
+        fault_tolerant_mode=fault_tolerant_mode)
+    num_elements = 2 * multiprocessing.cpu_count() + 100
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    iterator = iter(ds)
+    # Read halfway through the dataset.
+    midpoint = num_elements // 2
+    for i in range(midpoint):
+      self.assertEqual(i, next(iterator).numpy())
+
+    # Stop the original worker and start a new one.
+    cluster.workers[0].restart(use_same_port=use_same_port)
+
+    # There may have been some elements prefetched from the first worker
+    # before it was stopped.
+    while True:
+      val = next(iterator).numpy()
+      if val == 0:
+        break
+
+    # The dataset starts over now that we read from the new worker.
+    # TODO(b/157086991): Iterate until end of sequence when we support
+    # detecting lost workers.
+    for i in range(1, num_elements // 2):
+      val = next(iterator).numpy()
+      self.assertEqual(i, val)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testChangeProcessingModeAfterRestart(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 100
+    range_dataset = dataset_ops.Dataset.range(num_elements)
+    ds = range_dataset.apply(
+        data_service_ops.distribute(
+            processing_mode=data_service_ops.ShardingPolicy.OFF,
+            service=cluster.dispatcher_address(),
+            job_name="test"))
+    iterator = iter(ds)
+    for i in range(num_elements // 2):
+      self.assertEqual(i, next(iterator).numpy())
+    cluster.restart_dispatcher()
+    ds = range_dataset.apply(
+        data_service_ops.distribute(
+            processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+            service=cluster.dispatcher_address(),
+            job_name="test"))
+    with self.assertRaisesOpError(
+        "Tried to create job with name test, but found an existing job with "
+        "different parameters"):
+      next(iter(ds)).numpy()
+
+  @combinations.generate(
+      combinations.times(
+          test_base.eager_only_combinations(),
+          combinations.combine(work_dir=[TMP_WORK_DIR, NO_WORK_DIR])))
+  def testDistributeLargeGraphThenRegisterWorker(self, work_dir):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=0, work_dir=work_dir, fault_tolerant_mode=False)
+    # Larger than default OSS grpc message size limit of 4MB.
+    tensor = array_ops.ones((2, 1000, 1000), dtype=dtypes.float32)
+    ds = dataset_ops.Dataset.from_tensors(tensor)
+    ds = self.make_distributed_dataset(ds, cluster)
+    it = iter(ds)
+    cluster.add_worker()
+    self.assertAllEqual(next(it), tensor)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/gpu_test.py

@@ -0,0 +1,77 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import config
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.platform import test
+
+
+class TfDataServiceGpuTest(
+    data_service_test_base.TestBase,
+    parameterized.TestCase,
+):
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              pinned=[False, True],
+              data_transfer_protocol=["grpc", "local"],
+              compression=[False, True],
+          ),
+      )
+  )
+  def test_pinned(self, pinned, data_transfer_protocol, compression):
+    cpus = config.list_logical_devices("CPU")
+    gpus = config.list_logical_devices("GPU")
+    if not gpus:
+      self.skipTest("GPUs must be present to check GPU-pinnedness.")
+
+    num_elements = 10
+    cluster = self.make_test_cluster(num_workers=1)
+
+    with ops.device_v2(cpus[0].name):
+      ds = self.make_distributed_range_dataset(
+          num_elements=num_elements,
+          cluster=cluster,
+          data_transfer_protocol=data_transfer_protocol,
+          compression=("AUTO" if compression else None),
+      )
+
+    with ops.device_v2(gpus[0].name):
+      ds = ds.map(gen_experimental_dataset_ops.check_pinned)
+
+    options = options_lib.Options()
+    options.experimental_service.pinned = pinned
+    ds = ds.with_options(options)
+
+    if not pinned or data_transfer_protocol != "grpc" or compression:
+      with self.assertRaisesRegex(errors.InvalidArgumentError, "not pinned"):
+        self.assertDatasetProduces(ds, list(range(num_elements)))
+      return
+
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/local_workers_test.py

@@ -0,0 +1,441 @@
+# Copyright 2021 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 tf.data service with local and remote workers."""
+
+import multiprocessing
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import multi_process_cluster
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+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
+
+
+class LocalWorkersTest(data_service_test_base.TestBase, parameterized.TestCase):
+  """Tests reading from local workers if `target_workers` is `local`."""
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testOneLocalWorker(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=1, num_remote_workers=5)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="local")
+    self.assertDatasetProduces(ds, list(range(num_elements)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testLocalWorkers(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+    self.assertDatasetProduces(
+        ds,
+        num_local_workers * list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testRepeatedDataset(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    num_repetitions = 5
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+    ds = ds.repeat(num_repetitions)
+    self.assertDatasetProduces(
+        ds,
+        expected_output=num_local_workers * num_repetitions *
+        list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testPrefetchingDataset(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+    ds = ds.prefetch(10)
+    self.assertDatasetProduces(
+        ds,
+        expected_output=num_local_workers * list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testMultipleEpochs(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+    for _ in range(10):
+      self.assertDatasetProduces(
+          ds,
+          num_local_workers * list(range(num_elements)),
+          assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testDynamicSharding(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 100
+    ds = self.make_distributed_range_dataset(
+        num_elements,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC,
+        target_workers="LOCAL")
+    self.assertDatasetProduces(
+        ds, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleConsumers(self):
+    num_local_workers, num_remote_workers = 1, 3
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    # Because the elements in datasets are prefetched one per
+    # CPU core, a static number here may be excessively large
+    # for small numbers of CPU cores, or too small for high
+    # CPU core count machines, or probably both.
+    # In this case the below formula should satisfy both needs.
+    num_elements = 50 + (multiprocessing.cpu_count() * 2)
+    num_consumers = 8
+    iterators = []
+    for _ in range(num_consumers):
+      dataset = self.make_distributed_range_dataset(
+          num_elements, cluster, job_name="shared_job")
+      iterators.append(self.getNext(dataset))
+
+    results = []
+    for _ in range(10):
+      for it in iterators:
+        results.append(self.evaluate(it()))
+    for it in iterators:
+      results.extend(self.getIteratorOutput(it))
+
+    self.assertCountEqual(results, (num_local_workers + num_remote_workers) *
+                          list(range(num_elements)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testEmptyDataset(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 0
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+    self.assertDatasetProduces(ds, [])
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[0, 3], num_remote_workers=[1, 3])))
+  def testNonLocalRead(self, num_local_workers, num_remote_workers):
+    """This test ensures the remote workers are running and producing data."""
+
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(num_elements, cluster)
+    num_workers = num_local_workers + num_remote_workers
+    self.assertDatasetProduces(
+        ds, num_workers * list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoLocalWorker(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=0, num_remote_workers=3)
+    num_elements = 10
+    ds = self.make_distributed_range_dataset(
+        num_elements, cluster, target_workers="LOCAL")
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Local reads require local tf.data workers, but no local worker is "
+        "found."):
+      self.getDatasetOutput(ds)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInconsistentTargetWorkers(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=3, num_remote_workers=3)
+    ds = dataset_ops.Dataset.range(10)
+    datasets = [
+        self.make_distributed_dataset(
+            ds, cluster, job_name="test_job", target_workers=target_workers)
+        for target_workers in ["AUTO", "ANY", "LOCAL"]
+    ]
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "but found an existing job with different parameters: "
+        "Existing target workers: <AUTO>"):
+      for dataset in datasets:
+        self.getDatasetOutput(dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAnonymousJobWithDifferentTargetWorkers(self):
+    num_local_workers, num_remote_workers = (3, 3)
+    cluster = multi_process_cluster.MultiProcessCluster(num_local_workers,
+                                                        num_remote_workers)
+    num_elements = 10
+    ds = dataset_ops.Dataset.range(num_elements)
+    datasets = {
+        target_workers: self.make_distributed_dataset(
+            ds, cluster, target_workers=target_workers)
+        for target_workers in ["AUTO", "ANY", "LOCAL"]
+    }
+
+    num_workers = num_local_workers + num_remote_workers
+    self.assertDatasetProduces(
+        datasets["AUTO"],
+        num_workers * list(range(num_elements)),
+        assert_items_equal=True)
+    self.assertDatasetProduces(
+        datasets["ANY"],
+        num_workers * list(range(num_elements)),
+        assert_items_equal=True)
+    self.assertDatasetProduces(
+        datasets["LOCAL"],
+        num_local_workers * list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCoordinatedRead(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=3, num_remote_workers=3)
+    ds = dataset_ops.Dataset.range(10).repeat()
+    ds = self.make_distributed_dataset(
+        ds,
+        cluster,
+        job_name="test_job",
+        consumer_index=0,
+        num_consumers=3,
+        target_workers="LOCAL")
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Coordinated reads require non-local workers"):
+      self.getDatasetOutput(ds)
+
+
+class LocalTaskGarbageCollectTest(data_service_test_base.TestBase,
+                                  parameterized.TestCase):
+  """Tests garbage collecting unused local worker tasks.
+
+  The user typically creates an iterator in each epoch. This should delete the
+  previous iterator and releases the resources of it.
+  """
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testMultipleEpochs(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_epochs, num_steps = 5, 5
+    dataset = self._make_distributed_infinite_range_dataset(cluster)
+    for _ in range(num_epochs):
+      # For each iteration, the previous iterator is garbage collected.
+      get_next = self.getNext(dataset)
+      for i in range(num_steps):
+        self.assertEqual(self.evaluate(get_next()), i)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testMultipleEpochsSharedJob(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_epochs, num_steps = 5, 5
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    for _ in range(num_epochs):
+      # For each iteration, the previous iterator is garbage collected.
+      get_next = self.getNext(dataset)
+      for i in range(num_steps):
+        self.assertEqual(self.evaluate(get_next()), i)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_remote_workers=[0, 3], job_name=[None, "shared_job_name"])))
+  def testRepeatDistributedDataset(self, num_remote_workers, job_name):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    dataset = self.make_distributed_range_dataset(
+        10, cluster, job_name=job_name, target_workers="LOCAL")
+    dataset = dataset.repeat(3)
+    self.assertDatasetProduces(dataset, list(range(10)) * 3)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testReadFromDeletedTask(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_steps = 10
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    get_next = self.getNext(dataset)
+    for i in range(num_steps):
+      self.assertEqual(self.evaluate(get_next()), i)
+
+    # Re-creating the dataset resets the iterator index, so the second iterator
+    # reads from the same task as the first, which has been deleted.
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                "which has been deleted."):
+      get_next = self.getNext(dataset)
+      while True:
+        _ = self.evaluate(get_next())
+
+  @combinations.generate(
+      combinations.times(test_base.graph_only_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testReadFromDeletedTask_GraphMode(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_steps = 10
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    with self.session() as sess:
+      get_next = self.getNext(dataset)
+      for i in range(num_steps):
+        self.assertEqual(sess.run(get_next()), i)
+
+    # Re-creating the dataset resets the iterator index, so the second iterator
+    # reads from the same task as the first, which has been deleted.
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                "which has been deleted."):
+      with self.session() as sess:
+        get_next = self.getNext(dataset)
+        sess.run(get_next())
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testMultipleEpochs_WorkerRestart(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_steps = 10
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+
+    get_next = self.getNext(dataset)
+    for i in range(num_steps):
+      self.assertEqual(self.evaluate(get_next()), i)
+
+    # Verifies the worker re-creates the task after the iterator is deleted and
+    # the worker restarts.
+    del get_next
+    cluster.restart_local_workers()
+
+    get_next = self.getNext(dataset)
+    for i in range(num_steps):
+      self.assertEqual(self.evaluate(get_next()), i)
+
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testMultipleEpochs_DispatcherRestart(self, num_remote_workers):
+    num_local_workers = 1
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+
+    num_steps = 10
+    dataset = self._make_distributed_infinite_range_dataset(
+        cluster, job_name="shared_job_name")
+    get_next = self.getNext(dataset)
+    for i in range(num_steps):
+      self.assertEqual(self.evaluate(get_next()), i)
+
+    # Verifies the worker re-creates the task after the iterator is deleted and
+    # the dispatcher restarts.
+    del get_next
+    cluster.restart_dispatcher()
+
+    get_next = self.getNext(dataset)
+    for i in range(num_steps):
+      self.assertEqual(self.evaluate(get_next()), i)
+
+  def _make_distributed_infinite_range_dataset(self, cluster, job_name=None):
+    dataset = dataset_ops.Dataset.range(1000000).repeat()
+    return self.make_distributed_dataset(
+        dataset,
+        cluster=cluster,
+        job_name=job_name,
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        target_workers="LOCAL")
+
+
+if __name__ == "__main__":
+  multi_process_cluster.test_main()

tensorflow/python/data/experimental/kernel_tests/service/metadata_test.py

@@ -0,0 +1,228 @@
+# Copyright 2021 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 service metadata."""
+
+import functools
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distribute
+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_util
+from tensorflow.python.ops import script_ops
+from tensorflow.python.platform import test
+
+
+def _cardinality_test_combinations():
+  """Generate test combinations for data service cardinality tests.
+
+  We test only V2 combinations for the infinite and 0 cases because the `map`
+  transformation for compression makes the cardinality unknown in TF1.
+
+  Returns:
+    test combinations.
+  """
+
+  def _reduce_cases_to_combinations(result, case):
+    name, dataset_fn, sharding_policy, expected_result = case
+    return result + combinations.combine(
+        dataset_fn=combinations.NamedObject(name, dataset_fn),
+        sharding_policy=sharding_policy,
+        expected_result=expected_result)
+
+  def _cases_to_combinations(cases):
+    return functools.reduce(_reduce_cases_to_combinations, cases, [])
+
+  def _infinite_dataset_with_hint_shard():
+    return (dataset_ops.Dataset.range(10).shard(distribute.SHARD_HINT,
+                                                distribute.SHARD_HINT).repeat())
+
+  def _empty_dataset_with_hint_shard():
+    return (dataset_ops.Dataset.range(0).shard(distribute.SHARD_HINT,
+                                               distribute.SHARD_HINT))
+
+  v2_only_cases = [
+      ("NoShardingInfinite", lambda: dataset_ops.Dataset.range(10).repeat(),
+       data_service_ops.ShardingPolicy.OFF, dataset_ops.INFINITE),
+      ("DynamicShardingInfinite", lambda: dataset_ops.Dataset.range(5).repeat(),
+       data_service_ops.ShardingPolicy.DYNAMIC, dataset_ops.INFINITE),
+      ("DataShardingInfinite", lambda: dataset_ops.Dataset.range(10).repeat(),
+       data_service_ops.ShardingPolicy.DATA, dataset_ops.INFINITE),
+      ("NoShardingZero", lambda: dataset_ops.Dataset.range(0),
+       data_service_ops.ShardingPolicy.OFF, 0),
+      ("DynamicShardingZero", lambda: dataset_ops.Dataset.range(0),
+       data_service_ops.ShardingPolicy.DYNAMIC, 0),
+      ("DataShardingZero", lambda: dataset_ops.Dataset.range(0),
+       data_service_ops.ShardingPolicy.DATA, 0),
+      ("FileOrDataShardingZero", lambda: dataset_ops.Dataset.range(0),
+       data_service_ops.ShardingPolicy.FILE_OR_DATA, 0),
+      ("HintShardingZero", _empty_dataset_with_hint_shard,
+       data_service_ops.ShardingPolicy.HINT, dataset_ops.UNKNOWN),
+  ]
+  v1_and_v2_cases = [
+      ("Finite", lambda: dataset_ops.Dataset.range(10),
+       data_service_ops.ShardingPolicy.OFF, dataset_ops.UNKNOWN),
+      ("FileOrDataShardingUnknown",
+       lambda: dataset_ops.Dataset.range(10).repeat(),
+       data_service_ops.ShardingPolicy.FILE_OR_DATA, dataset_ops.UNKNOWN),
+      ("HintShardingUnknown", _infinite_dataset_with_hint_shard,
+       data_service_ops.ShardingPolicy.HINT, dataset_ops.UNKNOWN),
+  ]
+
+  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 v2_only_combinations + v1_and_v2_combinations
+
+
+class DataServiceMetadataTest(data_service_test_base.TestBase,
+                              parameterized.TestCase):
+  """Tests propagating data service metadata through tf.data service."""
+
+  @combinations.generate(_cardinality_test_combinations())
+  def testCardinality(self, dataset_fn, sharding_policy, expected_result):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    dataset = dataset_fn()
+    dataset = self.make_distributed_dataset(
+        dataset, cluster=cluster, processing_mode=sharding_policy)
+    self.assertEqual(self.evaluate(dataset.cardinality()), expected_result)
+
+  @combinations.generate(_cardinality_test_combinations())
+  def testFromDatasetIdCardinality(self, dataset_fn, sharding_policy,
+                                   expected_result):
+    cluster = data_service_test_base.TestCluster(num_workers=2)
+    dataset = dataset_fn()
+    dataset_id = data_service_ops.register_dataset(
+        cluster.dispatcher.target, dataset=dataset)
+    dataset = data_service_ops.from_dataset_id(
+        processing_mode=sharding_policy,
+        service=cluster.dispatcher.target,
+        dataset_id=dataset_id,
+        element_spec=dataset.element_spec)
+    self.assertEqual(self.evaluate(dataset.cardinality()), expected_result)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testFromDatasetIdDoesntRequireElementSpec(self):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1,
+        work_dir=data_service_test_base.NO_WORK_DIR,
+        fault_tolerant_mode=False,
+        data_transfer_protocol="grpc")
+    num_elements = 10
+    dataset = dataset_ops.Dataset.range(num_elements)
+
+    dataset_id = data_service_ops.register_dataset(cluster.dispatcher_address(),
+                                                   dataset)
+    dataset = data_service_ops.from_dataset_id(
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        service=cluster.dispatcher_address(),
+        dataset_id=dataset_id)
+    self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testElementSpecGraphMode(self):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1,
+        work_dir=data_service_test_base.NO_WORK_DIR,
+        fault_tolerant_mode=False)
+    num_elements = 10
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset_id = data_service_ops.register_dataset(cluster.dispatcher_address(),
+                                                   dataset)
+    with self.assertRaisesRegex(
+        ValueError, "In graph mode `element_spec` must be provided manually."):
+      _ = data_service_ops.from_dataset_id(
+          processing_mode=data_service_ops.ShardingPolicy.OFF,
+          service=cluster.dispatcher_address(),
+          dataset_id=dataset_id)
+
+  @combinations.generate(test_base.eager_only_combinations())
+  def testElementSpecMixedMode(self):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1,
+        work_dir=data_service_test_base.NO_WORK_DIR,
+        fault_tolerant_mode=False)
+    num_elements = 10
+    dataset = dataset_ops.Dataset.range(num_elements)
+
+    @def_function.function
+    def get_dataset_id():
+      return data_service_ops.register_dataset(cluster.dispatcher_address(),
+                                               dataset)
+
+    dataset_id = get_dataset_id()
+    dataset_id_val = tensor_util.constant_value(dataset_id)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        f"Failed to fetch element spec for dataset id {dataset_id_val} from "
+        "tf.data service. If the dataset was registered in graph mode or "
+        "inside a tf.function, the `element_spec` must be specified as an "
+        "argument to `from_dataset_id`."):
+      dataset = data_service_ops.from_dataset_id(
+          processing_mode=data_service_ops.ShardingPolicy.OFF,
+          service=cluster.dispatcher_address(),
+          dataset_id=dataset_id)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(compression=[None, "AUTO"])))
+  def testFromDatasetIdOmitsCompression(self, compression):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1, data_transfer_protocol="grpc")
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        list("abcdefghijklmnopqrstuvwxyz"))
+    def to_upper(x):
+      return script_ops.numpy_function(
+          func=lambda x: x.decode("utf-8").upper(), inp=[x], Tout=dtypes.string)
+    dataset = dataset.map(to_upper, num_parallel_calls=dataset_ops.AUTOTUNE)
+    dataset_id = data_service_ops.register_dataset(
+        cluster.dispatcher.target, dataset=dataset, compression=compression)
+    dataset = data_service_ops.from_dataset_id(
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        service=cluster.dispatcher.target,
+        dataset_id=dataset_id,
+        element_spec=dataset.element_spec)
+    self.assertDatasetProduces(dataset, list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"))
+
+  # Eager-only as querying `element_spec` is only supported in the eager mode.
+  @combinations.generate(
+      combinations.times(test_base.eager_only_combinations(),
+                         combinations.combine(compression=[None, "AUTO"])))
+  def testFromDatasetIdOmitsElementSpecAndCompression(self, compression):
+    cluster = data_service_test_base.TestCluster(
+        num_workers=1, data_transfer_protocol="grpc")
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"))
+    dataset_id = data_service_ops.register_dataset(
+        cluster.dispatcher.target, dataset=dataset, compression=compression)
+    dataset = data_service_ops.from_dataset_id(
+        processing_mode=data_service_ops.ShardingPolicy.OFF,
+        service=cluster.dispatcher.target,
+        dataset_id=dataset_id)
+    self.assertDatasetProduces(dataset, list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/multi_device_test.py

@@ -0,0 +1,84 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Multi-device tests for tf.data service ops."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+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.framework import ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.platform import test
+
+
+class MultiDeviceTest(data_service_test_base.TestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(MultiDeviceTest, self).setUp()
+    self._devices = self.configureDevicesForMultiDeviceTest(2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadDatasetOnDifferentDevices(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    num_elements = 10
+    with ops.device(self._devices[0]):
+      dataset = dataset_ops.Dataset.range(num_elements)
+      element_spec = dataset.element_spec
+      dataset_id = data_service_ops.register_dataset(
+          cluster.dispatcher_address(), dataset)
+      dataset = data_service_ops.from_dataset_id(
+          processing_mode=data_service_ops.ShardingPolicy.OFF,
+          service=cluster.dispatcher_address(),
+          dataset_id=dataset_id,
+          element_spec=element_spec)
+      self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+    with ops.device(self._devices[1]):
+      dataset = data_service_ops.from_dataset_id(
+          processing_mode=data_service_ops.ShardingPolicy.OFF,
+          service=cluster.dispatcher_address(),
+          dataset_id=dataset_id,
+          element_spec=dataset.element_spec)
+      self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testResourceOnWrongDevice(self):
+    cluster = data_service_test_base.TestCluster(num_workers=1)
+    with ops.device(self._devices[0]):
+      initializer = self.lookupTableInitializer("keyvaluetensor", [10, 11])
+      table = lookup_ops.StaticHashTable(initializer, -1)
+      self.evaluate(lookup_ops.tables_initializer())
+      dataset = dataset_ops.Dataset.range(3)
+      dataset = dataset.map(table.lookup)
+      dataset = self.make_distributed_dataset(dataset, cluster)
+      self.assertDatasetProduces(
+          dataset, [10, 11, -1], requires_initialization=True)
+
+    with ops.device(self._devices[1]):
+      dataset = dataset_ops.Dataset.range(3)
+      dataset = dataset.map(table.lookup)
+      with self.assertRaisesRegex(
+          errors.FailedPreconditionError,
+          "Serialization error while trying to register a dataset"):
+        dataset = self.make_distributed_dataset(dataset, cluster)
+        self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/kernel_tests/service/multi_process_cluster.py

@@ -0,0 +1,165 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""tf.data service test-cluster with local and remote workers."""
+
+import tempfile
+
+from tensorflow.core.protobuf import data_service_pb2
+from tensorflow.core.protobuf import service_config_pb2
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.service import server_lib
+from tensorflow.python.distribute import multi_process_lib
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+_WORKER_SHUTDOWN_QUIET_PERIOD_MS = 100
+
+
+# pylint: disable=protected-access
+class _RemoteWorkerProcess(multi_process_lib.Process):
+  """Runs a worker server in a new process to simulate a remote worker."""
+
+  def __init__(self, dispatcher_address, port, worker_tags, pipe_writer):
+    super(_RemoteWorkerProcess, self).__init__()
+    self._dispatcher_address = dispatcher_address
+    self._port = port
+    self._worker_tags = worker_tags
+    self._pipe_writer = pipe_writer
+
+  def run(self):
+    self.start_worker()
+
+  def start_worker(self):
+    self._worker = data_service_test_base.TestWorker(
+        self._dispatcher_address,
+        _WORKER_SHUTDOWN_QUIET_PERIOD_MS,
+        port=self._port,
+        worker_tags=self._worker_tags)
+    self._worker.start()
+    self._pipe_writer.send(self._worker.worker_address())
+    self._worker.join()
+
+
+class MultiProcessCluster:
+  """tf.data service cluster with local and remote workers.
+
+  Represents a cluster with a dispatcher, `num_local_workers` local workers, and
+  `num_remote_workers` remote workers. Remote workers run in separate processes.
+  This is useful to test reading from local in-process workers. For example:
+
+  ```
+  cluster = multi_process_cluster.MultiProcessCluster(
+      num_local_workers=1, num_remote_workers=3)
+  num_elements = 10
+  dataset = self.make_distributed_range_dataset(
+      num_elements, cluster, target_workers="LOCAL")
+  self.assertDatasetProduces(dataset, list(range(num_elements)))
+  ```
+  """
+
+  def __init__(self,
+               num_local_workers,
+               num_remote_workers,
+               worker_tags=None,
+               worker_addresses=None,
+               deployment_mode=data_service_pb2.DEPLOYMENT_MODE_COLOCATED):
+    self._work_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
+    self._deployment_mode = deployment_mode
+    self._start_dispatcher(worker_addresses)
+    self._start_local_workers(num_local_workers, worker_tags)
+    self._start_remote_workers(num_remote_workers, worker_tags)
+
+  def _start_dispatcher(self, worker_addresses, port=0):
+    if port == 0:
+      port = test_util.pick_unused_port()
+    self._dispatcher = server_lib.DispatchServer(
+        service_config_pb2.DispatcherConfig(
+            port=port,
+            protocol="grpc",
+            work_dir=self._work_dir,
+            fault_tolerant_mode=True,
+            worker_addresses=worker_addresses,
+            deployment_mode=self._deployment_mode),
+        start=True)
+
+  def _start_local_workers(self, num_workers, worker_tags=None):
+    self._local_workers = []
+    for _ in range(num_workers):
+      self.start_local_worker(worker_tags)
+
+  def _start_remote_workers(self, num_workers, worker_tags=None):
+    # List of (worker address, remote worker process) tuples.
+    self._remote_workers = []
+    for _ in range(num_workers):
+      self.start_remote_worker(worker_tags)
+
+  def start_local_worker(self, worker_tags=None):
+    worker = data_service_test_base.TestWorker(
+        self.dispatcher_address(),
+        _WORKER_SHUTDOWN_QUIET_PERIOD_MS,
+        port=test_util.pick_unused_port(),
+        worker_tags=worker_tags)
+    worker.start()
+    self._local_workers.append(worker)
+
+  def start_remote_worker(self, worker_tags=None):
+    """Runs a tf.data service worker in a remote process."""
+
+    pipe_reader, pipe_writer = multi_process_lib.multiprocessing.Pipe(
+        duplex=False)
+    worker_process = _RemoteWorkerProcess(
+        self.dispatcher_address(),
+        port=test_util.pick_unused_port(),
+        worker_tags=worker_tags,
+        pipe_writer=pipe_writer)
+    worker_process.start()
+    worker_address = pipe_reader.recv()
+    self._remote_workers.append((worker_address, worker_process))
+
+  def restart_dispatcher(self):
+    port = int(self.dispatcher_address().split(":")[1])
+    self._dispatcher._stop()
+    self._start_dispatcher(
+        worker_addresses=(self.local_worker_addresses() +
+                          self.remote_worker_addresses()),
+        port=port)
+
+  def restart_local_workers(self):
+    for worker in self._local_workers:
+      worker.restart()
+
+  def dispatcher_address(self):
+    return self._dispatcher._address
+
+  def local_worker_addresses(self):
+    return [worker.worker_address() for worker in self._local_workers]
+
+  def remote_worker_addresses(self):
+    return [worker_address for (worker_address, _) in self._remote_workers]
+
+  def _stop(self):
+    for worker in self._local_workers:
+      worker.stop()
+    for (_, worker_process) in self._remote_workers:
+      worker_process.kill()
+    self._dispatcher._stop()
+
+  def __del__(self):
+    self._stop()
+
+
+def test_main():
+  """Main function to be called within `__main__` of a test file."""
+  multi_process_lib.test_main()

tensorflow/python/data/experimental/kernel_tests/service/multi_process_cluster_test.py

@@ -0,0 +1,105 @@
+# Copyright 2021 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 tf.data service cluster with local and remote workers."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import data_service_pb2
+from tensorflow.python.data.experimental.kernel_tests.service import multi_process_cluster
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.ops import distribute
+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 test_ops
+from tensorflow.python.ops import math_ops
+
+
+class MultiProcessClusterTest(data_service_test_base.TestBase,
+                              parameterized.TestCase):
+  """Verifies the local and remote workers are running and producing data."""
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[0, 1, 3], num_remote_workers=[0, 1, 3])))
+  def testCluster(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers)
+    num_elements = 10
+    num_workers = num_local_workers + num_remote_workers
+    if num_workers == 0:
+      return
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    self.assertDatasetProduces(
+        dataset,
+        num_workers * list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDistributeNonblockingWithStuckWorkers(self):
+    num_workers = 6
+    # Avoids using local workers because it will stall the teardown
+    # while separate worker processes can be killed easily.
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=0,
+        num_remote_workers=num_workers,
+        worker_addresses=["localhost"] * num_workers,
+        deployment_mode=data_service_pb2.DEPLOYMENT_MODE_REMOTE,
+    )
+
+    num_elements = 10
+
+    def force_one_worker_to_stall_map(x):
+      # Simulates having a stuck worker.
+      if math_ops.equal(x, 0):
+        test_ops.sleep_op(sleep_seconds=10000)
+        return math_ops.cast(0, dtypes.int64)
+      else:
+        return x
+
+    dataset = dataset_ops.Dataset.range(num_elements, dtype=dtypes.int64)
+    dataset = dataset.shard(distribute.SHARD_HINT, distribute.SHARD_HINT)
+    dataset = dataset.map(force_one_worker_to_stall_map)
+    dataset = dataset.repeat()
+
+    dataset = self.make_distributed_dataset(
+        dataset,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.HINT,
+        # Makes sure only there is one client request at most.
+        max_outstanding_requests=1,
+    )
+
+    get_next = self.getNext(dataset, requires_initialization=False)
+
+    results = []
+    for _ in range(100):
+      results.append(self.evaluate(get_next()))
+
+    self.assertNotIn(
+        0,
+        results,
+        "The worker producing 0 should be sleeping so 0 should not show up in"
+        " the results.",
+    )
+
+
+if __name__ == "__main__":
+  multi_process_cluster.test_main()

tensorflow/python/data/experimental/kernel_tests/service/test_base.py

@@ -0,0 +1,456 @@
+# 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.
+# ==============================================================================
+"""Test base for tf.data service tests."""
+
+import os
+import shutil
+import tempfile
+
+from tensorflow.core.protobuf import service_config_pb2
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.experimental.service import server_lib
+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.ops import math_ops
+from tensorflow.python.platform import googletest
+
+# This will be resolved to a tmp directory by `start_dispatch_server`.
+TMP_WORK_DIR = "tmp_work_dir_placeholder"
+# `""` indicates not to use a work directory.
+NO_WORK_DIR = ""
+# We use a faster than normal heartbeat interval so that tests run faster.
+TEST_HEARTBEAT_INTERVAL_MS = 100
+TEST_DISPATCHER_TIMEOUT_MS = 5000
+TEST_WORKER_TIMEOUT_MS = 200
+TEST_JOB_GC_CHECK_INTERNAL_MS = 1000
+TEST_SNAPSHOT_MAX_CHUNK_SIZE_BYTES = 16 << 10  # 16 KB
+PROTOCOL = "grpc"
+
+
+def all_cluster_configurations():
+  with_work_dir = combinations.combine(
+      work_dir=TMP_WORK_DIR, fault_tolerant_mode=[True, False])
+  without_work_dir = combinations.combine(
+      work_dir=NO_WORK_DIR, fault_tolerant_mode=False)
+  return with_work_dir + without_work_dir
+
+
+def _make_worker(
+    dispatcher_address,
+    protocol,
+    data_transfer_protocol,
+    shutdown_quiet_period_ms=0,
+    port=0,
+    worker_tags=None,
+    cross_trainer_cache_size_bytes=None,
+    snapshot_max_chunk_size_bytes=TEST_SNAPSHOT_MAX_CHUNK_SIZE_BYTES,
+):
+  """Creates a worker server."""
+  defaults = server_lib.WorkerConfig(dispatcher_address=dispatcher_address)
+  config_proto = service_config_pb2.WorkerConfig(
+      dispatcher_address=dispatcher_address,
+      worker_address=defaults.worker_address,
+      port=port,
+      protocol=protocol,
+      worker_tags=worker_tags,
+      heartbeat_interval_ms=TEST_HEARTBEAT_INTERVAL_MS,
+      dispatcher_timeout_ms=TEST_DISPATCHER_TIMEOUT_MS,
+      data_transfer_protocol=data_transfer_protocol,
+      data_transfer_address=defaults.data_transfer_address,
+      shutdown_quiet_period_ms=shutdown_quiet_period_ms,
+      cross_trainer_cache_size_bytes=cross_trainer_cache_size_bytes,
+      snapshot_max_chunk_size_bytes=snapshot_max_chunk_size_bytes,
+  )
+  return server_lib.WorkerServer(config_proto, start=False)
+
+
+# pylint: disable=protected-access
+class TestWorker:
+  """A tf.data service worker."""
+
+  def __init__(
+      self,
+      dispatcher_address,
+      shutdown_quiet_period_ms,
+      protocol=PROTOCOL,
+      data_transfer_protocol=None,
+      port=0,
+      worker_tags=None,
+      cross_trainer_cache_size_bytes=None,
+      snapshot_max_chunk_size_bytes=TEST_SNAPSHOT_MAX_CHUNK_SIZE_BYTES,
+  ):
+    self._dispatcher_address = dispatcher_address
+    self._shutdown_quiet_period_ms = shutdown_quiet_period_ms
+    self._server = _make_worker(
+        dispatcher_address,
+        protocol,
+        data_transfer_protocol,
+        shutdown_quiet_period_ms,
+        port=port,
+        worker_tags=worker_tags,
+        cross_trainer_cache_size_bytes=cross_trainer_cache_size_bytes,
+        snapshot_max_chunk_size_bytes=snapshot_max_chunk_size_bytes,
+    )
+    self._running = False
+    self._protocol = protocol
+    self._data_transfer_protocol = data_transfer_protocol
+
+  def stop(self):
+    self._server._stop()
+    self._running = False
+
+  def start(self):
+    self._server.start()
+    self._port = int(self._server._address.split(":")[1])
+    self._running = True
+
+  def restart(self, use_same_port=True):
+    """Restarts the worker, stopping it first if it is already running."""
+    if self._running:
+      self.stop()
+    port = 0
+    if use_same_port:
+      port = self._port
+    self._server = _make_worker(self._dispatcher_address,
+                                self._protocol,
+                                self._data_transfer_protocol,
+                                self._shutdown_quiet_period_ms, port)
+    self._server.start()
+    self._port = int(self._server._address.split(":")[1])
+    self._running = True
+
+  def join(self):
+    self._server.join()
+
+  def num_tasks(self):
+    return self._server._num_tasks()
+
+  def snapshot_task_progresses(self):
+    return self._server._snapshot_task_progresses()
+
+  def worker_address(self):
+    return self._server._address
+
+
+class TestCluster:
+  """Test tf.data service cluster."""
+
+  def __init__(
+      self,
+      num_workers,
+      dispatcher_port=0,
+      work_dir=TMP_WORK_DIR,
+      fault_tolerant_mode=True,
+      job_gc_check_interval_ms=TEST_JOB_GC_CHECK_INTERNAL_MS,
+      job_gc_timeout_ms=None,
+      worker_timeout_ms=TEST_WORKER_TIMEOUT_MS,
+      worker_shutdown_quiet_period_ms=0,
+      snapshot_max_chunk_size_bytes=TEST_SNAPSHOT_MAX_CHUNK_SIZE_BYTES,
+      worker_max_concurrent_snapshots=0,
+      start=True,
+      protocol=PROTOCOL,
+      data_transfer_protocol=None,
+  ):
+    """Creates a tf.data service test cluster.
+
+    Args:
+      num_workers: The number of workers to initially add to the cluster.
+      dispatcher_port: The port to use for the dispatcher.
+      work_dir: The work directory to use for the dispatcher. If set to
+        `TMP_WORK_DIR`, the cluster will create a new temporary directory to use
+        as the work directory. If set to `NO_WORK_DIR`, no work directory will
+        be used.
+      fault_tolerant_mode: Whether the dispatcher should write its state to a
+        journal so that it can recover from restarts.
+      job_gc_check_interval_ms: How often the dispatcher should scan through to
+        delete old and unused jobs, in milliseconds.
+      job_gc_timeout_ms: How long a job needs to be unused before it becomes a
+        candidate for garbage collection, in milliseconds.
+      worker_timeout_ms: How long to wait for a worker to heartbeat before
+        considering it missing, in milliseconds.
+      worker_shutdown_quiet_period_ms: When shutting down a worker, how long to
+        wait for the gRPC server to process the final requests.
+      snapshot_max_chunk_size_bytes: The maximum size of a distributed snapshot
+        chunk file.
+      worker_max_concurrent_snapshots: The maximum number of snapshots a worker
+        can concurrently process.
+      start: Whether to immediately start the servers in the cluster. If
+        `False`, the servers can be started later by calling
+        `start_dispatcher()` and `start_workers()`.
+      protocol: The protocol to use for communicating with the tf.data service,
+        e.g. "grpc".
+      data_transfer_protocol: (Optional.) The protocol to use for transferring
+        data with the tf.data service.
+    """
+    if work_dir == TMP_WORK_DIR:
+      work_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
+    self._worker_shutdown_quiet_period_ms = worker_shutdown_quiet_period_ms
+    self._snapshot_max_chunk_size_bytes = snapshot_max_chunk_size_bytes
+    self._protocol = protocol
+    self._data_transfer_protocol = data_transfer_protocol
+    self._job_gc_check_interval_ms = job_gc_check_interval_ms
+    self._job_gc_timeout_ms = job_gc_timeout_ms
+    self._worker_timeout_ms = worker_timeout_ms
+    self._worker_max_concurrent_snapshots = worker_max_concurrent_snapshots
+    self.dispatcher = server_lib.DispatchServer(
+        server_lib.DispatcherConfig(
+            port=dispatcher_port,
+            work_dir=work_dir,
+            protocol=protocol,
+            fault_tolerant_mode=fault_tolerant_mode,
+            job_gc_check_interval_ms=job_gc_check_interval_ms,
+            job_gc_timeout_ms=job_gc_timeout_ms,
+            worker_timeout_ms=worker_timeout_ms,
+            worker_max_concurrent_snapshots=worker_max_concurrent_snapshots,
+        ),
+        start=start,
+    )
+
+    self.workers = []
+    for _ in range(num_workers):
+      self.add_worker(start=start)
+
+  def dispatcher_address(self):
+    return self.dispatcher.target.split("://")[1]
+
+  def add_worker(self, start=True):
+    worker = TestWorker(
+        self.dispatcher_address(),
+        self._worker_shutdown_quiet_period_ms,
+        self._protocol,
+        self._data_transfer_protocol,
+        snapshot_max_chunk_size_bytes=self._snapshot_max_chunk_size_bytes,
+    )
+    if start:
+      worker.start()
+    self.workers.append(worker)
+
+  def start_dispatcher(self):
+    self.dispatcher.start()
+
+  def start_workers(self):
+    for worker in self.workers:
+      worker.start()
+
+  def stop_dispatcher(self):
+    # pylint: disable=protected-access
+    self.dispatcher._stop()
+
+  def restart_worker(self, index):
+    self.workers[index].restart()
+
+  def stop_worker(self, index):
+    self.workers[index].stop()
+
+  def stop_workers(self):
+    for worker in self.workers:
+      worker.stop()
+
+  # pylint: disable=protected-access
+  def restart_dispatcher(self):
+    """Stops `dispatcher` and creates a new dispatcher with the same port.
+
+    Restarting is supported only when the dispatcher is configured with
+    `fault_tolerant_mode=True`.
+    """
+    if not self.dispatcher._config.fault_tolerant_mode:
+      raise ValueError(
+          "Trying to restart the dispatcher without fault-tolerance.")
+    port = int(self.dispatcher_address().split(":")[1])
+    self.dispatcher._stop()
+    self.dispatcher = server_lib.DispatchServer(
+        server_lib.DispatcherConfig(
+            port=port,
+            work_dir=self.dispatcher._config.work_dir,
+            protocol=self._protocol,
+            fault_tolerant_mode=self.dispatcher._config.fault_tolerant_mode,
+            job_gc_check_interval_ms=self._job_gc_check_interval_ms,
+            job_gc_timeout_ms=self._job_gc_timeout_ms,
+            worker_timeout_ms=self._worker_timeout_ms,
+            worker_max_concurrent_snapshots=
+            self._worker_max_concurrent_snapshots,
+        )
+    )
+
+  def num_registered_workers(self):
+    return self.dispatcher._num_workers()
+
+  def num_tasks_on_workers(self):
+    return sum(worker.num_tasks() for worker in self.workers)
+
+  def snapshot_streams(self, path):
+    return self.dispatcher._snapshot_streams(path)
+
+  def __del__(self):
+    # Destroy workers before the dispatcher for clean shutdown.
+    self.workers.clear()
+    del self.dispatcher
+
+
+class TestBase(test_base.DatasetTestBase):
+  """Base class for tf.data service tests."""
+
+  def setUp(self):
+    self.default_data_transfer_protocol = None
+    self.default_compression = "AUTO"
+
+  def set_default_data_transfer_protocol(self, protocol):
+    self.default_data_transfer_protocol = protocol
+
+  def set_default_compression(self, compression):
+    self.default_compression = compression
+
+  def make_test_cluster(self, *args, **kwargs):
+    if "data_transfer_protocol" not in kwargs:
+      kwargs["data_transfer_protocol"] = self.default_data_transfer_protocol
+    return TestCluster(*args, **kwargs)
+
+  def make_distributed_dataset(self,
+                               dataset,
+                               cluster,
+                               processing_mode="parallel_epochs",
+                               **kwargs):
+    kwargs["task_refresh_interval_hint_ms"] = 20
+    if "data_transfer_protocol" not in kwargs:
+      kwargs["data_transfer_protocol"] = self.default_data_transfer_protocol
+    if "compression" not in kwargs:
+      kwargs["compression"] = self.default_compression
+
+    # pylint: disable=protected-access
+    return dataset.apply(
+        data_service_ops._distribute(
+            processing_mode,
+            cluster.dispatcher_address(),
+            **kwargs))
+
+  def make_distributed_range_dataset(self,
+                                     num_elements,
+                                     cluster,
+                                     **kwargs):
+    dataset = dataset_ops.Dataset.range(num_elements)
+    return self.make_distributed_dataset(dataset, cluster, **kwargs)
+
+  def make_coordinated_read_dataset(
+      self,
+      cluster,
+      num_consumers,
+      sharding_policy=data_service_ops.ShardingPolicy.OFF):
+    """Creates a dataset that performs coordinated reads.
+
+    The dataset simulates `num_consumers` consumers by using parallel
+    interleave to read with `num_consumers` threads, one for each consumer. The
+    nth element of the dataset is produced by consumer `n % num_consumers`.
+
+    The dataset executed on each worker will produce groups of `num_consumers`
+    sequentially increasing numbers. For example, if `num_consumers=3` a worker
+    dataset could produce [0, 1, 2, 9, 10, 11, 21, 22, 23]. This enables
+    `checkCoordinatedReadGroups` below to assess whether the values received in
+    each step came from the same group.
+
+    Args:
+      cluster: A tf.data service `TestCluster`.
+      num_consumers: The number of consumers to simulate.
+      sharding_policy: The sharding policy to use. Currently only OFF and
+        DYNAMIC are supported.
+
+    Returns:
+      A dataset that simulates reading with `num_consumers` consumers.
+    """
+    if sharding_policy not in [
+        data_service_ops.ShardingPolicy.OFF,
+        data_service_ops.ShardingPolicy.DYNAMIC
+    ]:
+      raise ValueError(f"Unsupported sharding policy: {sharding_policy}")
+    # Start from 0 so that we can detect when a new worker is added with
+    # ShardingPolicy.OFF.
+    ds = dataset_ops.Dataset.from_tensors(math_ops.cast(0, dtypes.int64))
+    ds = ds.concatenate(dataset_ops.Dataset.random())
+    # Ensure that all elements in the same group are consecutive.
+    def make_group(x):
+      # Avoid overflowing an int64 in (x+1)*num_consumers below.
+      x = x % (2**32)
+      return dataset_ops.Dataset.range(x*num_consumers, (x+1)*num_consumers)
+    ds = ds.flat_map(make_group)
+    consumers = []
+    for consumer_index in range(num_consumers):
+      consumers.append(
+          self.make_distributed_dataset(
+              ds,
+              cluster,
+              job_name="test",
+              processing_mode=sharding_policy,
+              consumer_index=consumer_index,
+              num_consumers=num_consumers))
+    # Use parallel interleave to read from consumers in parallel.
+    ds = dataset_ops.Dataset.from_tensor_slices(consumers)
+    ds = ds.interleave(
+        lambda x: x,
+        cycle_length=num_consumers,
+        num_parallel_calls=num_consumers)
+    return ds
+
+  def checkCoordinatedReadGroups(self, results, num_consumers):
+    """Validates results from a `make_coordinted_read_dataset` dataset.
+
+    Each group of `num_consumers` results should be consecutive, indicating that
+    they were produced by the same worker.
+
+    Args:
+      results: The elements produced by the dataset.
+      num_consumers: The number of consumers.
+    """
+    groups = [
+        results[start:start + num_consumers]
+        for start in range(0, len(results), num_consumers)
+    ]
+    incorrect_groups = []
+    for group in groups:
+      # Check that each group of `num_consumers` results are consecutive.
+      for offset in range(1, len(group)):
+        if group[0] + offset != group[offset]:
+          incorrect_groups.append(group)
+          break
+    self.assertEmpty(
+        incorrect_groups,
+        "Incorrect groups: {}.\nAll groups: {}".format(incorrect_groups,
+                                                       groups))
+
+  def read(self, get_next, results, count):
+    for _ in range(count):
+      results.append(self.evaluate(get_next()))
+
+
+class TempDir:
+  """Temporary directory for unit testing."""
+
+  def __init__(self):
+    temp_dir = tempfile.mkdtemp(dir=googletest.GetTempDir())
+    self._path = os.path.join(
+        tempfile.mkdtemp(dir=temp_dir), "tf_data_snapshot")
+
+  @property
+  def full_path(self) -> str:
+    return self._path
+
+  def __fspath__(self) -> str:
+    return self._path
+
+  def __del__(self):
+    try:
+      shutil.rmtree(self.full_path)
+    except FileNotFoundError:
+      pass

tensorflow/python/data/experimental/kernel_tests/service/worker_tags_test.py

@@ -0,0 +1,210 @@
+# Copyright 2021 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 tf.data service reading from workers with specific tags."""
+
+import time
+import multiprocessing
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.kernel_tests.service import multi_process_cluster
+from tensorflow.python.data.experimental.kernel_tests.service import test_base as data_service_test_base
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.framework import combinations
+
+_COLOCATED_WORKER_TAG = "COLOCATED"
+
+
+class WorkerTagsTest(data_service_test_base.TestBase, parameterized.TestCase):
+  """Tests tf.data service reading from local or non-TPU workers.
+
+  When `target_workers` is "AUTO", tf.data service avoids cross-TPU reads, to
+  avoid RPCS and data serialization / deserialization, thus improving resource
+  utilization of TPU hosts.
+  """
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testReadFromLocalWorker(self, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=1,
+        num_remote_workers=num_remote_workers,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    # Only reads from the local worker.
+    self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testReadFromLocalAndNonTpuWorkers(self, num_local_workers,
+                                        num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    cluster.start_remote_worker(worker_tags=None)
+
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    # Reads from the local worker or non-colocated worker.
+    self.assertDatasetProduces(
+        dataset, (num_local_workers + 1) * list(range(num_elements)),
+        assert_items_equal=True)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         combinations.combine(num_remote_workers=[0, 3])))
+  def testLocalWorkerHasNoTag(self, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=0,
+        num_remote_workers=num_remote_workers,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    cluster.start_local_worker(worker_tags=None)
+
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    # Only reads from the local worker.
+    self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[0, 3])))
+  def testReadFromLocalAndNonTpuWorkers_DynamicSharding(self, num_local_workers,
+                                                        num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=3,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    cluster.start_remote_worker(worker_tags=None)
+
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(
+        num_elements,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.DYNAMIC)
+    self.assertDatasetProduces(
+        dataset, list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadFromLocalWorker_StaticSharding(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=1,
+        num_remote_workers=3,
+        worker_addresses=["localhost:%port%"] * 5,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    cluster.start_remote_worker(worker_tags=None)
+
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(
+        num_elements,
+        cluster,
+        processing_mode=data_service_ops.ShardingPolicy.FILE_OR_DATA)
+
+    # Static sharding will only read from the local worker.
+    self.assertDatasetProduces(dataset, list(range(0, num_elements, 5)))
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[1, 3])))
+  def testCoordinatedRead(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+    num_consumers = 4
+    dataset = self.make_coordinated_read_dataset(cluster, num_consumers)
+    get_next = self.getNext(dataset, requires_initialization=True)
+    results = [self.evaluate(get_next()) for _ in range(200)]
+    self.checkCoordinatedReadGroups(results, num_consumers)
+
+  @combinations.generate(
+      combinations.times(
+          test_base.default_test_combinations(),
+          combinations.combine(
+              num_local_workers=[1, 3], num_remote_workers=[1, 3])))
+  def testAddRemoteWorkersMidJob(self, num_local_workers, num_remote_workers):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=num_local_workers,
+        num_remote_workers=num_remote_workers,
+        worker_tags=[_COLOCATED_WORKER_TAG])
+
+    # num_elements needs to be bigger than (100 + <cpu core count>), the extra
+    # 100 is just a bit of margin. The CPU core count is involved as
+    # elements are prefetched, one element per CPU core.
+    num_elements = 200 + multiprocessing.cpu_count()
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    get_next = self.getNext(dataset)
+    results = [self.evaluate(get_next()) for _ in range(100)]
+
+    # Will only read from the two non-TPU workers.
+    cluster.start_remote_worker(worker_tags=None)
+    cluster.start_remote_worker(worker_tags=[_COLOCATED_WORKER_TAG])
+    cluster.start_remote_worker(worker_tags=None)
+    cluster.start_remote_worker(worker_tags=[_COLOCATED_WORKER_TAG])
+    expect_num_workers_to_read = num_local_workers + 2
+
+    # Wait for the new worker to register with the dispatcher.
+    while cluster._dispatcher._num_workers() < (num_local_workers +
+                                                num_remote_workers + 4):
+      time.sleep(10 / 1000)  # 10ms
+
+    results += self.getIteratorOutput(get_next)
+    self.assertCountEqual(
+        results, expect_num_workers_to_read * list(range(num_elements)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleTags(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=1,
+        num_remote_workers=3,
+        worker_tags=[_COLOCATED_WORKER_TAG, "COLOCATED_2", "COLOCATED_3"])
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    # Only reads from the local worker.
+    self.assertDatasetProduces(dataset, list(range(num_elements)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testUnusedTags(self):
+    cluster = multi_process_cluster.MultiProcessCluster(
+        num_local_workers=1,
+        num_remote_workers=3,
+        worker_tags=["Unused tag 1", "Unused tag 2", "Unused tag 3"])
+    num_elements = 100
+    dataset = self.make_distributed_range_dataset(num_elements, cluster)
+    # The tags don't have an effect. tf.data service will read from all workers.
+    self.assertDatasetProduces(
+        dataset, 4 * list(range(num_elements)), assert_items_equal=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidTag(self):
+    with self.assertRaisesRegex(RuntimeError, "Worker tags cannot be empty."):
+      _ = multi_process_cluster.MultiProcessCluster(
+          num_local_workers=1,
+          num_remote_workers=3,
+          worker_tags=["", _COLOCATED_WORKER_TAG])
+
+
+if __name__ == "__main__":
+  multi_process_cluster.test_main()

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

@@ -0,0 +1,167 @@
+# 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.shuffle_and_repeat()`."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.experimental.ops import shuffle_ops
+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.framework import random_seed
+from tensorflow.python.platform import test
+
+
+class ShuffleAndRepeatTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _build_ds(self, seed, count=5, num_elements=20):
+    return dataset_ops.Dataset.range(num_elements).apply(
+        shuffle_ops.shuffle_and_repeat(buffer_size=5, count=count, seed=seed))
+
+  def _gen_outputs(self, ds_fn, num_outputs, verify_exhausted=True):
+    get_next = self.getNext(ds_fn())
+    outputs = []
+    for _ in range(num_outputs):
+      outputs.append(self.evaluate(get_next()))
+    if verify_exhausted:
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(get_next())
+    return outputs
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCorrectOutput(self):
+    output = self._gen_outputs(lambda: self._build_ds(10), 100)
+    self.assertSequenceEqual(
+        sorted(output), sorted(
+            np.array([range(20) for _ in range(5)]).flatten()))
+    for i in range(5):
+      self.assertSequenceEqual(sorted(output[i * 20:(i + 1) * 20]), range(20))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testReshuffling(self):
+    # Check that the output orders of different epochs are indeed different.
+    output = self._gen_outputs(lambda: self._build_ds(10), 100)
+    for i in range(4):
+      epoch1 = output[i * 20:(i + 1) * 20]
+      epoch2 = output[(i + 1) * 20:(i + 2) * 20]
+      self.assertNotEqual(epoch1, epoch2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSameOrderForSameSeeds(self):
+    output1 = self._gen_outputs(lambda: self._build_ds(10), 100)
+    output2 = self._gen_outputs(lambda: self._build_ds(10), 100)
+    self.assertEqual(output1, output2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDifferentOrderForDifferentSeeds(self):
+    output1 = self._gen_outputs(lambda: self._build_ds(10), 100)
+    output2 = self._gen_outputs(lambda: self._build_ds(20), 100)
+    self.assertNotEqual(output1, output2)
+    self.assertCountEqual(output1, output2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCountNone(self):
+    output1 = self._gen_outputs(
+        lambda: self._build_ds(10, count=None), 100, verify_exhausted=False)
+    output2 = self._gen_outputs(
+        lambda: self._build_ds(20, count=None), 100, verify_exhausted=False)
+    self.assertNotEqual(output1, output2)
+    self.assertCountEqual(output1, output2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCountMinusOne(self):
+    output1 = self._gen_outputs(
+        lambda: self._build_ds(10, count=-1), 100, verify_exhausted=False)
+    output2 = self._gen_outputs(
+        lambda: self._build_ds(20, count=-1), 100, verify_exhausted=False)
+    self.assertNotEqual(output1, output2)
+    self.assertCountEqual(output1, output2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteOutputs(self):
+    # Asserting the iterator is exhausted after producing 100 items should fail.
+    with self.assertRaises(AssertionError):
+      self._gen_outputs(lambda: self._build_ds(10, count=None), 100)
+    with self.assertRaises(AssertionError):
+      self._gen_outputs(lambda: self._build_ds(10, count=-1), 100)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteEmpty(self):
+    with self.assertRaises(errors.OutOfRangeError):
+      self._gen_outputs(lambda: self._build_ds(10, count=None, num_elements=0),
+                        100)
+    with self.assertRaises(errors.OutOfRangeError):
+      self._gen_outputs(lambda: self._build_ds(10, count=-1, num_elements=0),
+                        100)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testLargeBufferSize(self):
+    ds = dataset_ops.Dataset.range(20).apply(
+        shuffle_ops.shuffle_and_repeat(buffer_size=21))
+    get_next = self.getNext(ds)
+    self.evaluate(get_next())
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testVeryLargeBufferSize(self):
+    num_epochs = 1000 * 1000
+    # Each element being shuffled and repeated has shape (100,). This will OOM
+    # or timeout if we actually load everything into the buffer.
+    ds = dataset_ops.Dataset.range(500).batch(100).apply(
+        shuffle_ops.shuffle_and_repeat(
+            buffer_size=5 * num_epochs, count=num_epochs))
+    # Verify two epochs worth of output.
+    output = self._gen_outputs(lambda: ds, 2 * 5, verify_exhausted=False)
+    for i in range(2):
+      sorted_epoch = sorted(
+          output[i * 5:(i + 1) * 5], key=lambda batch: batch[0])
+      self.assertAllEqual(sorted_epoch, np.arange(500).reshape([5, 100]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRerandomizeOnReplicate(self):
+    random_seed.set_random_seed(None)
+    # When no seeds are fixed, each instantiation of the dataset should
+    # produce elements in a different order.
+    num_epochs = 2
+    num_elements = 100
+    ds = dataset_ops.Dataset.range(num_elements).apply(
+        shuffle_ops.shuffle_and_repeat(
+            buffer_size=num_elements, count=num_epochs))
+
+    shuffle_1 = self.getDatasetOutput(ds)
+    ds = self.graphRoundTrip(ds)
+    shuffle_2 = self.getDatasetOutput(ds)
+
+    self.assertCountEqual(shuffle_1, shuffle_2)
+    self.assertNotEqual(shuffle_1, shuffle_2)
+
+
+class ShuffleAndRepeatCheckpointTest(checkpoint_test_base.CheckpointTestBase,
+                                     parameterized.TestCase):
+
+  def _build_ds(self, seed):
+    return dataset_ops.Dataset.range(20).apply(
+        shuffle_ops.shuffle_and_repeat(buffer_size=5, count=5, seed=seed))
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    verify_fn(self, lambda: self._build_ds(10), num_outputs=100)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,78 @@
+# 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.experimental.sleep()`."""
+import time
+
+from absl.testing import parameterized
+
+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.framework import combinations
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class SleepTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSleep(self):
+    self.skipTest("b/123597912")
+    sleep_microseconds = 100
+    dataset = dataset_ops.Dataset.range(10).apply(
+        testing.sleep(sleep_microseconds))
+    next_element = self.getNext(dataset)
+    start_time = time.time()
+    for i in range(10):
+      self.assertEqual(i, self.evaluate(next_element()))
+    end_time = time.time()
+    self.assertGreater(end_time - start_time, (10 * sleep_microseconds) / 1e6)
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element())
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testSleepCancellation(self):
+    sleep_microseconds = int(1e6) * 1000
+    ds = dataset_ops.Dataset.range(1)
+    ds = ds.apply(testing.sleep(sleep_microseconds))
+    ds = ds.prefetch(1)
+    get_next = self.getNext(ds, requires_initialization=True)
+
+    with self.cached_session() as sess:
+      thread = self.checkedThread(self.assert_op_cancelled, args=(get_next(),))
+      thread.start()
+      time.sleep(0.2)
+      sess.close()
+      thread.join()
+
+  @combinations.generate(combinations.combine(tf_api_version=1, mode="graph"))
+  def testSleepBackgroundCancellation(self):
+    ds = dataset_ops.Dataset.range(1)
+
+    sleep_microseconds = int(1e6) * 1000
+    ds_sleep = dataset_ops.Dataset.range(1)
+    ds_sleep = ds.apply(testing.sleep(sleep_microseconds))
+
+    ds = ds.concatenate(ds_sleep)
+    ds = ds.prefetch(1)
+
+    get_next = self.getNext(ds, requires_initialization=True)
+
+    with self.cached_session():
+      self.assertEqual(self.evaluate(get_next()), 0)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,613 @@
+# 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.SqlDataset`."""
+
+import os
+
+from absl.testing import parameterized
+
+import sqlite3
+
+from tensorflow.python.data.experimental.ops import readers
+from tensorflow.python.data.kernel_tests import checkpoint_test_base
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class SqlDatasetTestBase(test_base.DatasetTestBase):
+  """Base class for setting up and testing SqlDataset."""
+
+  def _createSqlDataset(self,
+                        query,
+                        output_types,
+                        driver_name="sqlite",
+                        num_repeats=1):
+    dataset = readers.SqlDataset(driver_name, self.data_source_name, query,
+                                 output_types).repeat(num_repeats)
+    return dataset
+
+  def setUp(self):
+    super(SqlDatasetTestBase, self).setUp()
+    self.data_source_name = os.path.join(test.get_temp_dir(), "tftest.sqlite")
+
+    conn = sqlite3.connect(self.data_source_name)
+    c = conn.cursor()
+    c.execute("DROP TABLE IF EXISTS students")
+    c.execute("DROP TABLE IF EXISTS people")
+    c.execute("DROP TABLE IF EXISTS townspeople")
+    c.execute("DROP TABLE IF EXISTS data")
+    c.execute(
+        "CREATE TABLE IF NOT EXISTS students (id INTEGER NOT NULL PRIMARY KEY, "
+        "first_name VARCHAR(100), last_name VARCHAR(100), motto VARCHAR(100), "
+        "school_id VARCHAR(100), favorite_nonsense_word VARCHAR(100), "
+        "desk_number INTEGER, income INTEGER, favorite_number INTEGER, "
+        "favorite_big_number INTEGER, favorite_negative_number INTEGER, "
+        "favorite_medium_sized_number INTEGER, brownie_points INTEGER, "
+        "account_balance INTEGER, registration_complete INTEGER)")
+    c.executemany(
+        "INSERT INTO students (first_name, last_name, motto, school_id, "
+        "favorite_nonsense_word, desk_number, income, favorite_number, "
+        "favorite_big_number, favorite_negative_number, "
+        "favorite_medium_sized_number, brownie_points, account_balance, "
+        "registration_complete) "
+        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+        [("John", "Doe", "Hi!", "123", "n\0nsense", 9, 0, 2147483647,
+          9223372036854775807, -2, 32767, 0, 0, 1),
+         ("Jane", "Moe", "Hi again!", "1000", "nonsense\0", 127, -20000,
+          -2147483648, -9223372036854775808, -128, -32768, 255, 65535, 0)])
+    c.execute(
+        "CREATE TABLE IF NOT EXISTS people (id INTEGER NOT NULL PRIMARY KEY, "
+        "first_name VARCHAR(100), last_name VARCHAR(100), state VARCHAR(100))")
+    c.executemany(
+        "INSERT INTO PEOPLE (first_name, last_name, state) VALUES (?, ?, ?)",
+        [("Benjamin", "Franklin", "Pennsylvania"), ("John", "Doe",
+                                                    "California")])
+    c.execute(
+        "CREATE TABLE IF NOT EXISTS townspeople (id INTEGER NOT NULL PRIMARY "
+        "KEY, first_name VARCHAR(100), last_name VARCHAR(100), victories "
+        "FLOAT, accolades FLOAT, triumphs FLOAT)")
+    c.executemany(
+        "INSERT INTO townspeople (first_name, last_name, victories, "
+        "accolades, triumphs) VALUES (?, ?, ?, ?, ?)",
+        [("George", "Washington", 20.00,
+          1331241.321342132321324589798264627463827647382647382643874,
+          9007199254740991.0),
+         ("John", "Adams", -19.95,
+          1331241321342132321324589798264627463827647382647382643874.0,
+          9007199254740992.0)])
+    c.execute("CREATE TABLE IF NOT EXISTS data (col1 INTEGER)")
+    c.executemany("INSERT INTO DATA VALUES (?)", [(0,), (1,), (2,)])
+    conn.commit()
+    conn.close()
+
+
+class SqlDatasetTest(SqlDatasetTestBase, parameterized.TestCase):
+
+  # Test that SqlDataset can read from a database table.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSet(self):
+    for _ in range(2):  # Run twice to verify statelessness of db operations.
+      dataset = self._createSqlDataset(
+          query="SELECT first_name, last_name, motto FROM students "
+          "ORDER BY first_name DESC",
+          output_types=(dtypes.string, dtypes.string, dtypes.string),
+          num_repeats=2)
+      self.assertDatasetProduces(
+          dataset,
+          expected_output=[(b"John", b"Doe", b"Hi!"),
+                           (b"Jane", b"Moe", b"Hi again!")] * 2,
+          num_test_iterations=2)
+
+  # Test that SqlDataset works on a join query.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetJoinQuery(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT students.first_name, state, motto FROM students "
+            "INNER JOIN people "
+            "ON students.first_name = people.first_name "
+            "AND students.last_name = people.last_name",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+
+    self.assertEqual((b"John", b"California", b"Hi!"),
+                     self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that SqlDataset can read a database entry with a null-terminator
+  # in the middle of the text and place the entry in a `string` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetNullTerminator(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, favorite_nonsense_word "
+            "FROM students ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+
+    self.assertEqual((b"John", b"Doe", b"n\0nsense"), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", b"Moe", b"nonsense\0"),
+                     self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that SqlDataset works when used on two different queries.
+  # Because the output types of the dataset must be determined at graph-creation
+  # time, the two queries must have the same number and types of columns.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetReuseSqlDataset(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, motto FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    self.assertEqual((b"John", b"Doe", b"Hi!"), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", b"Moe", b"Hi again!"), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, state FROM people "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    self.assertEqual((b"John", b"Doe", b"California"),
+                     self.evaluate(get_next()))
+    self.assertEqual((b"Benjamin", b"Franklin", b"Pennsylvania"),
+                     self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that an `OutOfRangeError` is raised on the first call to
+  # `get_next_str_only` if result set is empty.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadEmptyResultSet(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, motto FROM students "
+            "WHERE first_name = 'Nonexistent'",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that an error is raised when `driver_name` is invalid.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetWithInvalidDriverName(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      dataset = self._createSqlDataset(
+          driver_name="sqlfake",
+          query="SELECT first_name, last_name, motto FROM students "
+          "ORDER BY first_name DESC",
+          output_types=(dtypes.string, dtypes.string, dtypes.string))
+      self.assertDatasetProduces(dataset, expected_output=[])
+
+  # Test that an error is raised when a column name in `query` is nonexistent
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetWithInvalidColumnName(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, fake_column FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    with self.assertRaises(errors.UnknownError):
+      self.evaluate(get_next())
+
+  # Test that an error is raised when there is a syntax error in `query`.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetOfQueryWithSyntaxError(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELEmispellECT first_name, last_name, motto FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    with self.assertRaises(errors.UnknownError):
+      self.evaluate(get_next())
+
+  # Test that an error is raised when the number of columns in `query`
+  # does not match the length of `, output_types`.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetWithMismatchBetweenColumnsAndOutputTypes(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+
+  # Test that no results are returned when `query` is an insert query rather
+  # than a select query. In particular, the error refers to the number of
+  # output types passed to the op not matching the number of columns in the
+  # result set of the query (namely, 0 for an insert statement.)
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetOfInsertQuery(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="INSERT INTO students (first_name, last_name, motto) "
+            "VALUES ('Foo', 'Bar', 'Baz'), ('Fizz', 'Buzz', 'Fizzbuzz')",
+            output_types=(dtypes.string, dtypes.string, dtypes.string)))
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table and
+  # place it in an `int8` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt8(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int8)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a negative or 0-valued integer from a
+  # SQLite database table and place it in an `int8` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt8NegativeAndZero(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, income, favorite_negative_number "
+            "FROM students "
+            "WHERE first_name = 'John' ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int8, dtypes.int8)))
+    self.assertEqual((b"John", 0, -2), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a large (positive or negative) integer from
+  # a SQLite database table and place it in an `int8` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt8MaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT desk_number, favorite_negative_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.int8, dtypes.int8)))
+    self.assertEqual((9, -2), self.evaluate(get_next()))
+    # Max and min values of int8
+    self.assertEqual((127, -128), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table and
+  # place it in an `int16` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt16(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int16)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a negative or 0-valued integer from a
+  # SQLite database table and place it in an `int16` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt16NegativeAndZero(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, income, favorite_negative_number "
+            "FROM students "
+            "WHERE first_name = 'John' ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int16, dtypes.int16)))
+    self.assertEqual((b"John", 0, -2), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a large (positive or negative) integer from
+  # a SQLite database table and place it in an `int16` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt16MaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, favorite_medium_sized_number "
+            "FROM students ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int16)))
+    # Max value of int16
+    self.assertEqual((b"John", 32767), self.evaluate(get_next()))
+    # Min value of int16
+    self.assertEqual((b"Jane", -32768), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table and
+  # place it in an `int32` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt32(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int32)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+
+  # Test that `SqlDataset` can read a negative or 0-valued integer from a
+  # SQLite database table and place it in an `int32` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt32NegativeAndZero(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, income FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int32)))
+    self.assertEqual((b"John", 0), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", -20000), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a large (positive or negative) integer from
+  # a SQLite database table and place it in an `int32` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt32MaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, favorite_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int32)))
+    # Max value of int32
+    self.assertEqual((b"John", 2147483647), self.evaluate(get_next()))
+    # Min value of int32
+    self.assertEqual((b"Jane", -2147483648), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a numeric `varchar` from a SQLite database
+  # table and place it in an `int32` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt32VarCharColumnAsInt(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, school_id FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int32)))
+    self.assertEqual((b"John", 123), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 1000), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table
+  # and place it in an `int64` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt64(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int64)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a negative or 0-valued integer from a
+  # SQLite database table and place it in an `int64` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt64NegativeAndZero(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, income FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int64)))
+    self.assertEqual((b"John", 0), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", -20000), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a large (positive or negative) integer from
+  # a SQLite database table and place it in an `int64` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetInt64MaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, favorite_big_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.int64)))
+    # Max value of int64
+    self.assertEqual((b"John", 9223372036854775807), self.evaluate(get_next()))
+    # Min value of int64
+    self.assertEqual((b"Jane", -9223372036854775808), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table and
+  # place it in a `uint8` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetUInt8(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.uint8)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read the minimum and maximum uint8 values from a
+  # SQLite database table and place them in `uint8` tensors.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetUInt8MinAndMaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, brownie_points FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.uint8)))
+    # Min value of uint8
+    self.assertEqual((b"John", 0), self.evaluate(get_next()))
+    # Max value of uint8
+    self.assertEqual((b"Jane", 255), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer from a SQLite database table
+  # and place it in a `uint16` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetUInt16(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, desk_number FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.uint16)))
+    self.assertEqual((b"John", 9), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", 127), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read the minimum and maximum uint16 values from a
+  # SQLite database table and place them in `uint16` tensors.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetUInt16MinAndMaxValues(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, account_balance FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.uint16)))
+    # Min value of uint16
+    self.assertEqual((b"John", 0), self.evaluate(get_next()))
+    # Max value of uint16
+    self.assertEqual((b"Jane", 65535), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a 0-valued and 1-valued integer from a
+  # SQLite database table and place them as `True` and `False` respectively
+  # in `bool` tensors.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetBool(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, registration_complete FROM students "
+            "ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.bool)))
+    self.assertEqual((b"John", True), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", False), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read an integer that is not 0-valued or 1-valued
+  # from a SQLite database table and place it as `True` in a `bool` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetBoolNotZeroOrOne(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, favorite_medium_sized_number "
+            "FROM students ORDER BY first_name DESC",
+            output_types=(dtypes.string, dtypes.bool)))
+    self.assertEqual((b"John", True), self.evaluate(get_next()))
+    self.assertEqual((b"Jane", True), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a float from a SQLite database table
+  # and place it in a `float64` tensor.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetFloat64(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, victories FROM townspeople "
+            "ORDER BY first_name",
+            output_types=(dtypes.string, dtypes.string, dtypes.float64)))
+    self.assertEqual((b"George", b"Washington", 20.0),
+                     self.evaluate(get_next()))
+    self.assertEqual((b"John", b"Adams", -19.95), self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a float from a SQLite database table beyond
+  # the precision of 64-bit IEEE, without throwing an error. Test that
+  # `SqlDataset` identifies such a value as equal to itself.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetFloat64OverlyPrecise(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, accolades FROM townspeople "
+            "ORDER BY first_name",
+            output_types=(dtypes.string, dtypes.string, dtypes.float64)))
+    self.assertEqual(
+        (b"George", b"Washington",
+         1331241.321342132321324589798264627463827647382647382643874),
+        self.evaluate(get_next()))
+    self.assertEqual(
+        (b"John", b"Adams",
+         1331241321342132321324589798264627463827647382647382643874.0),
+        self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that `SqlDataset` can read a float from a SQLite database table,
+  # representing the largest integer representable as a 64-bit IEEE float
+  # such that the previous integer is also representable as a 64-bit IEEE float.
+  # Test that `SqlDataset` can distinguish these two numbers.
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetFloat64LargestConsecutiveWholeNumbersNotEqual(self):
+    get_next = self.getNext(
+        self._createSqlDataset(
+            query="SELECT first_name, last_name, triumphs FROM townspeople "
+            "ORDER BY first_name",
+            output_types=(dtypes.string, dtypes.string, dtypes.float64)))
+    self.assertNotEqual((b"George", b"Washington", 9007199254740992.0),
+                        self.evaluate(get_next()))
+    self.assertNotEqual((b"John", b"Adams", 9007199254740991.0),
+                        self.evaluate(get_next()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+  # Test that SqlDataset can stop correctly when combined with batch
+  @combinations.generate(test_base.default_test_combinations())
+  def testReadResultSetWithBatchStop(self):
+    dataset = self._createSqlDataset(
+        query="SELECT * FROM data", output_types=(dtypes.int32))
+    dataset = dataset.map(lambda x: array_ops.identity(x))
+    get_next = self.getNext(dataset.batch(2))
+    self.assertAllEqual(self.evaluate(get_next()), [0, 1])
+    self.assertAllEqual(self.evaluate(get_next()), [2])
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(get_next())
+
+
+class SqlDatasetCheckpointTest(SqlDatasetTestBase,
+                               checkpoint_test_base.CheckpointTestBase,
+                               parameterized.TestCase):
+
+  def _build_dataset(self, num_repeats):
+    data_source_name = os.path.join(test.get_temp_dir(), "tftest.sqlite")
+    driver_name = array_ops.placeholder_with_default(
+        array_ops.constant("sqlite", dtypes.string), shape=[])
+    query = ("SELECT first_name, last_name, motto FROM students ORDER BY "
+             "first_name DESC")
+    output_types = (dtypes.string, dtypes.string, dtypes.string)
+    return readers.SqlDataset(driver_name, data_source_name, query,
+                              output_types).repeat(num_repeats)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         checkpoint_test_base.default_test_combinations()))
+  def test(self, verify_fn):
+    num_repeats = 4
+    num_outputs = num_repeats * 2
+    verify_fn(self, lambda: self._build_dataset(num_repeats), num_outputs)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,145 @@
+# 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.experimental.TFRecordWriter`."""
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.experimental.ops import grouping
+from tensorflow.python.data.experimental.ops import writers
+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.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import dtypes
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.lib.io import tf_record
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class TFRecordWriterTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def setUp(self):
+    super(TFRecordWriterTest, self).setUp()
+    self._num_records = 8
+
+  def writer_fn(self, filename, compression_type=""):
+    input_dataset = readers.TFRecordDataset([filename], compression_type)
+    return writers.TFRecordWriter(self._outputFilename(),
+                                  compression_type).write(input_dataset)
+
+  def _record(self, i):
+    return compat.as_bytes("Record %d" % (i))
+
+  def _createFile(self, options=None):
+    filename = self._inputFilename()
+    writer = python_io.TFRecordWriter(filename, options)
+    for i in range(self._num_records):
+      writer.write(self._record(i))
+    writer.close()
+    return filename
+
+  def _inputFilename(self):
+    return os.path.join(self.get_temp_dir(), "tf_record.in.txt")
+
+  def _outputFilename(self):
+    return os.path.join(self.get_temp_dir(), "tf_record.out.txt")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWrite(self):
+    self.evaluate(self.writer_fn(self._createFile()))
+    for i, r in enumerate(tf_record.tf_record_iterator(self._outputFilename())):
+      self.assertAllEqual(self._record(i), r)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWriteZLIB(self):
+    options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.ZLIB)
+    self.evaluate(
+        self.writer_fn(self._createFile(options), compression_type="ZLIB"))
+    for i, r in enumerate(
+        tf_record.tf_record_iterator(self._outputFilename(), options=options)):
+      self.assertAllEqual(self._record(i), r)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWriteGZIP(self):
+    options = tf_record.TFRecordOptions(tf_record.TFRecordCompressionType.GZIP)
+    self.evaluate(
+        self.writer_fn(self._createFile(options), compression_type="GZIP"))
+    for i, r in enumerate(
+        tf_record.tf_record_iterator(self._outputFilename(), options=options)):
+      self.assertAllEqual(self._record(i), r)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFailDataset(self):
+    with self.assertRaises(TypeError):
+      writers.TFRecordWriter(self._outputFilename(), "").write("whoops")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFailDType(self):
+    input_dataset = dataset_ops.Dataset.from_tensors(10)
+    with self.assertRaises(TypeError):
+      writers.TFRecordWriter(self._outputFilename(), "").write(input_dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFailShape(self):
+    input_dataset = dataset_ops.Dataset.from_tensors([["hello"], ["world"]])
+    with self.assertRaises(TypeError):
+      writers.TFRecordWriter(self._outputFilename(), "").write(input_dataset)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testSideEffect(self):
+    def writer_fn():
+      input_dataset = readers.TFRecordDataset(self._createFile())
+      return writers.TFRecordWriter(self._outputFilename()).write(input_dataset)
+
+    @def_function.function
+    def fn():
+      _ = writer_fn()
+      return "hello"
+
+    self.assertEqual(self.evaluate(fn()), b"hello")
+    for i, r in enumerate(tf_record.tf_record_iterator(self._outputFilename())):
+      self.assertAllEqual(self._record(i), r)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShard(self):
+    filename = self._createFile()
+    dataset = readers.TFRecordDataset([filename])
+
+    def reduce_func(key, dataset):
+      shard_filename = string_ops.string_join(
+          [filename, string_ops.as_string(key)])
+      writer = writers.TFRecordWriter(shard_filename)
+      writer.write(dataset.map(lambda _, x: x))
+      return dataset_ops.Dataset.from_tensors(shard_filename)
+
+    dataset = dataset.enumerate()
+    dataset = dataset.apply(
+        grouping.group_by_window(lambda i, _: i % 2, reduce_func,
+                                 dtypes.int64.max))
+
+    get_next = self.getNext(dataset)
+    for i in range(2):
+      shard_filename = (filename + str(i)).encode()
+      self.assertEqual(self.evaluate(get_next()), shard_filename)
+      for j, r in enumerate(tf_record.tf_record_iterator(shard_filename)):
+        self.assertAllEqual(self._record(i + 2*j), r)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,47 @@
+# 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.experimental.{from,to}_variant()`."""
+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 VariantTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testRoundtripRange(self):
+    dataset = dataset_ops.Dataset.range(10)
+    variant = dataset_ops.to_variant(dataset)
+    dataset = dataset_ops.from_variant(variant,
+                                       dataset_ops.get_structure(dataset))
+    self.assertDatasetProduces(dataset, range(10))
+    self.assertEqual(self.evaluate(dataset.cardinality()), 10)
+
+  @combinations.generate(
+      combinations.combine(tf_api_version=[2], mode=["eager", "graph"]))
+  def testRoundtripMap(self):
+    dataset = dataset_ops.Dataset.range(10).map(lambda x: x * x)
+    variant = dataset_ops.to_variant(dataset)
+    dataset = dataset_ops.from_variant(variant,
+                                       dataset_ops.get_structure(dataset))
+    self.assertDatasetProduces(dataset, [x * x for x in range(10)])
+    self.assertEqual(self.evaluate(dataset.cardinality()), 10)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,269 @@
+# 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 `tf.data.Dataset.weighted_flat_map()`."""
+
+from typing import Callable
+import unittest
+
+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 weighted_flat_map_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.platform import test
+
+
+class WeightedFlatMapTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testWeightedFlatMap(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset3 = dataset_ops.Dataset.range(20, 30)
+    dataset = weighted_flat_map_op._weighted_flat_map(
+        [dataset1, dataset2, dataset3], np.asarray([1, 1, 2]))
+    self.assertDatasetProduces(
+        dataset, expected_output=list(range(5)) + list(range(10, 15)) +
+        list(range(20, 30)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInvalidCardinality(self):
+    dataset1 = dataset_ops.Dataset.range(100)
+    dataset2 = dataset_ops.Dataset.range(100, 200)
+    dataset3 = dataset_ops.Dataset.range(200, 210)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "Input.*needs to have at least."):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3], np.asarray([1, 1, 100]))
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteCardinality(self):
+    dataset1 = dataset_ops.Dataset.range(10).repeat()
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset3 = dataset_ops.Dataset.range(20, 30)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Cardinalities of the inputs must be known."):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3])
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testEmptyInputDatasets(self):
+    dataset1 = dataset_ops.Dataset.from_tensor_slices([])
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset3 = dataset_ops.Dataset.range(20, 30)
+    with self.assertRaisesRegex(
+        TypeError,
+        "Incompatible dataset elements"):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3])
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testZeroWeight(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset3 = dataset_ops.Dataset.range(20, 30)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "`weights` must be greater than 0.0"):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3], [0, 1.0, 1.0])
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInfiniteWeight(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "`weights` must be finite"
+    ):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2], [float("inf"), 1.0]
+      )
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMultipleInfiniteWeights(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "`weights` must be finite"
+    ):
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2], [float("inf"), float("inf")]
+      )
+      self.getDatasetOutput(dataset, requires_initialization=True)
+
+
+@unittest.skip("TODO(b/325112575): Fix incompatibility with batch dataset")
+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(10).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset2 = dataset_ops.Dataset.range(10, 20).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset3 = dataset_ops.Dataset.range(20, 30).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset = weighted_flat_map_op._weighted_flat_map(
+        [dataset1, dataset2, dataset3], np.asarray([0.25, 0.25, 0.5]))
+    dataset = global_shuffle_op._global_shuffle(dataset)
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertCountEqual(
+        output, list(range(5)) + list(range(10, 15)) + list(range(20, 30)))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffledInputs(self):
+    dataset1 = dataset_ops.Dataset.range(10).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset2 = dataset_ops.Dataset.range(10, 20).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset3 = dataset_ops.Dataset.range(20, 30).prefetch(
+        buffer_size=dataset_ops.AUTOTUNE)
+    dataset1 = global_shuffle_op._global_shuffle(dataset1, seed=42)
+    dataset2 = global_shuffle_op._global_shuffle(dataset2, seed=42)
+    dataset3 = global_shuffle_op._global_shuffle(dataset3, seed=42)
+    dataset = weighted_flat_map_op._weighted_flat_map(
+        [dataset1, dataset2, dataset3], np.asarray([0.25, 0.25, 0.5]))
+
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    # Verifies that the first 5 elements are from `dataset1` in a random order.
+    self.assertFalse(set(output[:5]).issubset(set(range(5))))
+    self.assertTrue(set(output[:5]).issubset(set(range(10))))
+    # Verifies that the second 5 elements are from `dataset2` in a random order.
+    self.assertFalse(set(output[5:10]).issubset(set(range(10, 15))))
+    self.assertTrue(set(output[5:10]).issubset(set(range(10, 20))))
+    # Verifies that the last 10 elements are from `dataset3` in a random order.
+    self.assertCountEqual(output[10:], range(20, 30))
+    self.assertNotEqual(output[10:], range(20, 30))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffledInputsAndOutput(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset3 = dataset_ops.Dataset.range(20, 30)
+    dataset1 = global_shuffle_op._global_shuffle(dataset1, seed=42)
+    dataset2 = global_shuffle_op._global_shuffle(dataset2, seed=42)
+    dataset3 = global_shuffle_op._global_shuffle(dataset3, seed=42)
+    dataset = weighted_flat_map_op._weighted_flat_map(
+        [dataset1, dataset2, dataset3], np.asarray([0.25, 0.25, 0.5]))
+    dataset = global_shuffle_op._global_shuffle(dataset, seed=42)
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    # Verifies that not all first 5 elements are from `dataset1`.
+    self.assertFalse(set(output[:5]).issubset(set(range(10))))
+    # Verifies that not all second 5 elements are from `dataset2`.
+    self.assertFalse(set(output[5:10]).issubset(set(range(10, 20))))
+    # Verifies that not all last 10 elements are from `dataset3`.
+    self.assertFalse(set(output[10:]).issubset(set(range(20, 30))))
+
+    sorted_output = sorted(output)
+    # Verifies that there are 5 elements from dataset1
+    self.assertTrue(set(sorted_output[:5]).issubset(set(range(10))))
+    # Verifies that there are 5 elements from dataset2
+    self.assertTrue(set(sorted_output[5:10]).issubset(set(range(10, 20))))
+    # Verifies that there are 10 elements from dataset3
+    self.assertTrue(set(sorted_output[10:]).issubset(set(range(20, 30))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testShuffledWithBatch(self):
+    dataset1 = dataset_ops.Dataset.range(10)
+    dataset2 = dataset_ops.Dataset.range(10, 20)
+    dataset1 = dataset1.batch(2, drop_remainder=True)
+    dataset2 = dataset2.batch(2, drop_remainder=True)
+    dataset = weighted_flat_map_op._weighted_flat_map(
+        [dataset1, dataset2], np.asarray([0.5, 0.5])
+    )
+    shuffled_dataset = global_shuffle_op._global_shuffle(dataset, seed=42)
+    shuffled_output = self.getDatasetOutput(
+        shuffled_dataset, requires_initialization=True
+    )
+    output = self.getDatasetOutput(dataset, requires_initialization=True)
+    self.assertLen(shuffled_output, len(output))
+
+
+@unittest.skip("TODO(b/325112575): Fix incompatibility with batch dataset")
+class WeightedFlatMapGlobalShuffleCheckpointTest(
+    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=[True, False])))
+  def testWeightedFlatMap(
+      self,
+      verify_fn: Callable[..., None],
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset1 = dataset_ops.Dataset.range(10)
+      dataset2 = dataset_ops.Dataset.range(10, 20)
+      dataset3 = dataset_ops.Dataset.range(20, 30)
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3], np.asarray([0.25, 0.25, 0.5]))
+      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=20)
+
+  @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 testGlobalshuffle(
+      self,
+      verify_fn: Callable[..., None],
+      reshuffle_each_iteration: bool,
+      symbolic_checkpoint: bool):
+
+    def _build_dataset() -> dataset_ops.Dataset:
+      dataset1 = dataset_ops.Dataset.range(10)
+      dataset2 = dataset_ops.Dataset.range(10, 20)
+      dataset3 = dataset_ops.Dataset.range(20, 30)
+      dataset = weighted_flat_map_op._weighted_flat_map(
+          [dataset1, dataset2, dataset3], np.asarray([0.25, 0.25, 0.5]))
+      dataset = global_shuffle_op._global_shuffle(
+          dataset, seed=42, 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=20,
+        assert_items_equal=reshuffle_each_iteration)
+
+
+if __name__ == "__main__":
+  test.main()

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

@@ -0,0 +1,72 @@
+# 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 wrapping / unwrapping dataset variants."""
+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.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.platform import test
+
+
+class WrapUnwrapTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  # TODO(b/182414964): After making options persistent across tf.function is
+  # enabled, the ModelDatasetOp and MaxIntraParallelismOp are no longer present
+  # in Python. As a result, the FinalizeDataset is placed on GPU because of
+  # colocation constraint on the iterator. It then requires a registered copy
+  # operation from CPU to GPU for RangeDataset that does not exist and the test
+  # fails. Fix this test and re-enable it.
+  def DISABLED_testBasic(self):
+    ds = dataset_ops.Dataset.range(100)
+    ds_variant = ds._variant_tensor  # pylint: disable=protected-access
+
+    wrapped_variant = gen_dataset_ops.wrap_dataset_variant(ds_variant)
+    unwrapped_variant = gen_dataset_ops.unwrap_dataset_variant(wrapped_variant)
+
+    variant_ds = dataset_ops._VariantDataset(unwrapped_variant,
+                                             ds.element_spec)
+    get_next = self.getNext(variant_ds, requires_initialization=True)
+    for i in range(100):
+      self.assertEqual(i, self.evaluate(get_next()))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testGPU(self):
+    ds = dataset_ops.Dataset.range(100)
+    ds_variant = ds._variant_tensor  # pylint: disable=protected-access
+    wrapped_variant = gen_dataset_ops.wrap_dataset_variant(ds_variant)
+
+    with ops.device("/gpu:0"):
+      gpu_wrapped_variant = array_ops.identity(wrapped_variant)
+
+    unwrapped_variant = gen_dataset_ops.unwrap_dataset_variant(
+        gpu_wrapped_variant)
+    variant_ds = dataset_ops._VariantDataset(unwrapped_variant,
+                                             ds.element_spec)
+    iterator = dataset_ops.make_initializable_iterator(variant_ds)
+    get_next = iterator.get_next()
+
+    with self.cached_session():
+      self.evaluate(iterator.initializer)
+      for i in range(100):
+        self.assertEqual(i, self.evaluate(get_next))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/ops/BUILD

@@ -0,0 +1,635 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "batching",
+    srcs = ["batching.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "cardinality",
+    srcs = ["cardinality.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "compression_ops",
+    srcs = ["compression_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "counter",
+    srcs = ["counter.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "data_service_ops",
+    srcs = [
+        "data_service_ops.py",
+    ],
+    lazy_imports = True,
+    strict_deps = True,
+    deps = [
+        ":compression_ops",
+        "//tensorflow/core/protobuf:for_core_protos_py_proto",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/data/experimental/service:_pywrap_server_lib",
+        "//tensorflow/python/data/experimental/service:_pywrap_utils_exp",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "data_service_ops_test",
+    srcs = ["data_service_ops_test.py"],
+    deps = [
+        ":data_service_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "distributed_save_op",
+    srcs = [
+        "distributed_save_op.py",
+    ],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "distribute",
+    srcs = [
+        "distribute.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/types:data",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "enumerate_ops",
+    srcs = ["enumerate_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "error_ops",
+    srcs = ["error_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "from_list",
+    srcs = ["from_list.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "get_single_element",
+    srcs = ["get_single_element.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/types:data",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "global_shuffle_op",
+    srcs = [
+        "global_shuffle_op.py",
+    ],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/util:random_seed",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "grouping",
+    srcs = ["grouping.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "index_flat_map_op",
+    srcs = [
+        "index_flat_map_op.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "interleave_ops",
+    srcs = ["interleave_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "io",
+    srcs = [
+        "io.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "iterator_ops",
+    srcs = [
+        "iterator_ops.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:proto_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "iterator_model_ops",
+    srcs = [
+        "iterator_model_ops.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "lookup_ops",
+    srcs = [
+        "lookup_ops.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":cardinality",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:lookup_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "map_defun",
+    srcs = ["map_defun.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "matching_files",
+    srcs = ["matching_files.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "pad_to_cardinality",
+    srcs = ["pad_to_cardinality.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "parsing_ops",
+    srcs = ["parsing_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "prefetching_ops",
+    srcs = ["prefetching_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:device",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:functional_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "random_access",
+    srcs = ["random_access.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "random_ops",
+    srcs = [
+        "random_ops.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "readers",
+    srcs = [
+        "readers.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":error_ops",
+        ":parsing_ops",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "resampling",
+    srcs = ["resampling.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "scan_ops",
+    srcs = ["scan_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "shuffle_ops",
+    srcs = [
+        "shuffle_ops.py",
+    ],
+    strict_deps = True,
+    visibility = [
+        # TODO(jsimsa): Remove non-default visibility exception when `index_shuffle` is made public.
+        "//third_party/py/tensorflow_datasets:__subpackages__",
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":random_access",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:random_seed",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateless_random_ops",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "snapshot",
+    srcs = [
+        "snapshot.py",
+    ],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "take_while_ops",
+    srcs = ["take_while_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "testing",
+    srcs = ["testing.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "unique",
+    srcs = [
+        "unique.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "weighted_flat_map_op",
+    srcs = [
+        "weighted_flat_map_op.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:random_seed",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+    ],
+)
+
+py_library(
+    name = "writers",
+    srcs = [
+        "writers.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/types:data",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "dataset_ops",
+    strict_deps = True,
+    deps = [
+        ":batching",
+        ":cardinality",
+        ":compression_ops",
+        ":counter",
+        ":data_service_ops",
+        ":distribute",
+        ":distributed_save_op",
+        ":enumerate_ops",
+        ":error_ops",
+        ":from_list",
+        ":get_single_element",
+        ":global_shuffle_op",
+        ":grouping",
+        ":interleave_ops",
+        ":io",
+        ":map_defun",
+        ":matching_files",
+        ":pad_to_cardinality",
+        ":prefetching_ops",
+        ":random_access",
+        ":readers",
+        ":resampling",
+        ":scan_ops",
+        ":shuffle_ops",
+        ":snapshot",
+        ":take_while_ops",
+        ":unique",
+        ":writers",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:debug_mode",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/ops:dataset_ops_gen",
+    ],
+)

tensorflow/python/data/experimental/ops/batching.py

@@ -0,0 +1,379 @@
+# 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.
+# ==============================================================================
+"""Batching dataset transformations."""
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import convert
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+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_util
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.dense_to_ragged_batch")
+@deprecation.deprecated(None, "Use `tf.data.Dataset.ragged_batch` instead.")
+def dense_to_ragged_batch(batch_size,
+                          drop_remainder=False,
+                          row_splits_dtype=dtypes.int64):
+  """A transformation that batches ragged elements into `tf.RaggedTensor`s.
+
+  This transformation combines multiple consecutive elements of the input
+  dataset into a single element.
+
+  Like `tf.data.Dataset.batch`, the components of the resulting element will
+  have an additional outer dimension, which will be `batch_size` (or
+  `N % batch_size` for the last element if `batch_size` does not divide the
+  number of input elements `N` evenly and `drop_remainder` is `False`). If
+  your program depends on the batches having the same outer dimension, you
+  should set the `drop_remainder` argument to `True` to prevent the smaller
+  batch from being produced.
+
+  Unlike `tf.data.Dataset.batch`, the input elements to be batched may have
+  different shapes:
+
+  *  If an input element is a `tf.Tensor` whose static `tf.TensorShape` is
+     fully defined, then it is batched as normal.
+  *  If an input element is a `tf.Tensor` whose static `tf.TensorShape` contains
+     one or more axes with unknown size (i.e., `shape[i]=None`), then the output
+     will contain a `tf.RaggedTensor` that is ragged up to any of such
+     dimensions.
+  *  If an input element is a `tf.RaggedTensor` or any other type, then it is
+     batched as normal.
+
+  Example:
+
+  >>> dataset = tf.data.Dataset.from_tensor_slices(np.arange(6))
+  >>> dataset = dataset.map(lambda x: tf.range(x))
+  >>> dataset.element_spec.shape
+  TensorShape([None])
+  >>> dataset = dataset.apply(
+  ...     tf.data.experimental.dense_to_ragged_batch(batch_size=2))
+  >>> for batch in dataset:
+  ...   print(batch)
+  <tf.RaggedTensor [[], [0]]>
+  <tf.RaggedTensor [[0, 1], [0, 1, 2]]>
+  <tf.RaggedTensor [[0, 1, 2, 3], [0, 1, 2, 3, 4]]>
+
+  Args:
+    batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+      consecutive elements of this dataset to combine in a single batch.
+    drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+      whether the last batch should be dropped in the case it has fewer than
+      `batch_size` elements; the default behavior is not to drop the smaller
+      batch.
+    row_splits_dtype: The dtype that should be used for the `row_splits` of any
+      new ragged tensors.  Existing `tf.RaggedTensor` elements do not have their
+      row_splits dtype changed.
+
+  Returns:
+    Dataset: A `Dataset`.
+  """
+  def _apply_fn(dataset):
+    return dataset.ragged_batch(batch_size, drop_remainder, row_splits_dtype)
+
+  return _apply_fn
+
+
+@tf_export("data.experimental.dense_to_sparse_batch")
+@deprecation.deprecated(None, "Use `tf.data.Dataset.sparse_batch` instead.")
+def dense_to_sparse_batch(batch_size, row_shape):
+  """A transformation that batches ragged elements into `tf.sparse.SparseTensor`s.
+
+  Like `Dataset.padded_batch()`, this transformation combines multiple
+  consecutive elements of the dataset, which might have different
+  shapes, into a single element. The resulting element has three
+  components (`indices`, `values`, and `dense_shape`), which
+  comprise a `tf.sparse.SparseTensor` that represents the same data. The
+  `row_shape` represents the dense shape of each row in the
+  resulting `tf.sparse.SparseTensor`, to which the effective batch size is
+  prepended. For example:
+
+  ```python
+  # NOTE: The following examples use `{ ... }` to represent the
+  # contents of a dataset.
+  a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
+
+  a.apply(tf.data.experimental.dense_to_sparse_batch(
+      batch_size=2, row_shape=[6])) ==
+  {
+      ([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1]],  # indices
+       ['a', 'b', 'c', 'a', 'b'],                 # values
+       [2, 6]),                                   # dense_shape
+      ([[0, 0], [0, 1], [0, 2], [0, 3]],
+       ['a', 'b', 'c', 'd'],
+       [1, 6])
+  }
+  ```
+
+  Args:
+    batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+      consecutive elements of this dataset to combine in a single batch.
+    row_shape: A `tf.TensorShape` or `tf.int64` vector tensor-like object
+      representing the equivalent dense shape of a row in the resulting
+      `tf.sparse.SparseTensor`. Each element of this dataset must have the same
+      rank as `row_shape`, and must have size less than or equal to `row_shape`
+      in each dimension.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+  """
+
+  def _apply_fn(dataset):
+    return dataset.sparse_batch(batch_size, row_shape)
+
+  return _apply_fn
+
+
+@deprecation.deprecated(None, "Use `tf.data.experimental.map_and_batch()")
+@tf_export(v1=["data.experimental.map_and_batch_with_legacy_function"])
+def map_and_batch_with_legacy_function(map_func,
+                                       batch_size,
+                                       num_parallel_batches=None,
+                                       drop_remainder=False,
+                                       num_parallel_calls=None):
+  """Fused implementation of `map` and `batch`.
+
+  NOTE: This is an escape hatch for existing uses of `map_and_batch` that do not
+  work with V2 functions. New uses are strongly discouraged and existing uses
+  should migrate to `map_and_batch` as this method will not be removed in V2.
+
+  Args:
+    map_func: A function mapping a nested structure of tensors to another
+      nested structure of tensors.
+    batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+      consecutive elements of this dataset to combine in a single batch.
+    num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
+      representing the number of batches to create in parallel. On one hand,
+      higher values can help mitigate the effect of stragglers. On the other
+      hand, higher values can increase contention if CPU is scarce.
+    drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+      whether the last batch should be dropped in case its size is smaller than
+      desired; the default behavior is not to drop the smaller batch.
+    num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
+      representing the number of elements to process in parallel. If not
+      specified, `batch_size * num_parallel_batches` elements will be processed
+      in parallel. If the value `tf.data.AUTOTUNE` is used, then
+      the number of parallel calls is set dynamically based on available CPU.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+
+  Raises:
+    ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
+      specified.
+  """
+
+  if num_parallel_batches is None and num_parallel_calls is None:
+    num_parallel_calls = batch_size
+  elif num_parallel_batches is not None and num_parallel_calls is None:
+    num_parallel_calls = batch_size * num_parallel_batches
+  elif num_parallel_batches is not None and num_parallel_calls is not None:
+    raise ValueError(
+        "`map_and_batch_with_legacy_function` allows only one of "
+        "`num_parallel_batches` and "
+        "`num_parallel_calls` to be set, but "
+        f"`num_parallel_batches` was set to {num_parallel_batches} "
+        f"and `num_parallel_calls` as set to {num_parallel_calls}.")
+
+  def _apply_fn(dataset):
+    return _MapAndBatchDataset(dataset, map_func, batch_size,
+                               num_parallel_calls, drop_remainder,
+                               use_legacy_function=True)
+
+  return _apply_fn
+
+
+@deprecation.deprecated(
+    None,
+    "Use `tf.data.Dataset.map(map_func, num_parallel_calls)` followed by "
+    "`tf.data.Dataset.batch(batch_size, drop_remainder)`. Static tf.data "
+    "optimizations will take care of using the fused implementation.")
+@tf_export("data.experimental.map_and_batch")
+def map_and_batch(map_func,
+                  batch_size,
+                  num_parallel_batches=None,
+                  drop_remainder=False,
+                  num_parallel_calls=None):
+  """Fused implementation of `map` and `batch`.
+
+  Maps `map_func` across `batch_size` consecutive elements of this dataset
+  and then combines them into a batch. Functionally, it is equivalent to `map`
+  followed by `batch`. This API is temporary and deprecated since input pipeline
+  optimization now fuses consecutive `map` and `batch` operations automatically.
+
+  Args:
+    map_func: A function mapping a nested structure of tensors to another
+      nested structure of tensors.
+    batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
+      consecutive elements of this dataset to combine in a single batch.
+    num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
+      representing the number of batches to create in parallel. On one hand,
+      higher values can help mitigate the effect of stragglers. On the other
+      hand, higher values can increase contention if CPU is scarce.
+    drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+      whether the last batch should be dropped in case its size is smaller than
+      desired; the default behavior is not to drop the smaller batch.
+    num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
+      representing the number of elements to process in parallel. If not
+      specified, `batch_size * num_parallel_batches` elements will be processed
+      in parallel. If the value `tf.data.AUTOTUNE` is used, then
+      the number of parallel calls is set dynamically based on available CPU.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+
+  Raises:
+    ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
+      specified.
+  """
+
+  if num_parallel_batches is None and num_parallel_calls is None:
+    num_parallel_calls = batch_size
+  elif num_parallel_batches is not None and num_parallel_calls is None:
+    num_parallel_calls = batch_size * num_parallel_batches
+  elif num_parallel_batches is not None and num_parallel_calls is not None:
+    raise ValueError(
+        "`map_and_batch` allows only one of `num_parallel_batches` and "
+        "`num_parallel_calls` to be set, but "
+        f"`num_parallel_batches` was set to {num_parallel_batches} "
+        f"and `num_parallel_calls` as set to {num_parallel_calls}.")
+
+  def _apply_fn(dataset):
+    return _MapAndBatchDataset(dataset, map_func, batch_size,
+                               num_parallel_calls, drop_remainder)
+
+  return _apply_fn
+
+
+@deprecation.deprecated(None, "Use `tf.data.Dataset.unbatch()`.")
+@tf_export("data.experimental.unbatch")
+def unbatch():
+  """Splits elements of a dataset into multiple elements on the batch dimension.
+
+  For example, if elements of the dataset are shaped `[B, a0, a1, ...]`,
+  where `B` may vary for each input element, then for each element in the
+  dataset, the unbatched dataset will contain `B` consecutive elements
+  of shape `[a0, a1, ...]`.
+
+  ```python
+  # NOTE: The following example uses `{ ... }` to represent the contents
+  # of a dataset.
+  a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }
+
+  a.unbatch() == {
+      'a', 'b', 'c', 'a', 'b', 'a', 'b', 'c', 'd'}
+  ```
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+  """
+
+  def _apply_fn(dataset):
+    return dataset.unbatch()
+
+  return _apply_fn
+
+
+class _DenseToSparseBatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that batches ragged dense elements into `tf.sparse.SparseTensor`s."""
+
+  def __init__(self, input_dataset, batch_size, row_shape):
+    """See `Dataset.dense_to_sparse_batch()` for more details."""
+    if not isinstance(
+        dataset_ops.get_legacy_output_types(input_dataset), dtypes.DType):
+      raise TypeError("`dense_to_sparse_batch` requires an input dataset whose "
+                      "elements have a single component, but the given dataset "
+                      "has the following component types: "
+                      f"{dataset_ops.get_legacy_output_types(input_dataset)}.")
+    self._input_dataset = input_dataset
+    self._batch_size = batch_size
+    self._row_shape = row_shape
+    self._element_spec = sparse_tensor.SparseTensorSpec(
+        tensor_shape.TensorShape([None]).concatenate(self._row_shape),
+        dataset_ops.get_legacy_output_types(input_dataset))
+
+    variant_tensor = ged_ops.dense_to_sparse_batch_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._batch_size,
+        row_shape=convert.partial_shape_to_tensor(self._row_shape),
+        **self._flat_structure)
+    super(_DenseToSparseBatchDataset, self).__init__(input_dataset,
+                                                     variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+class _MapAndBatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over a batch of elements."""
+
+  def __init__(self, input_dataset, map_func, batch_size, num_parallel_calls,
+               drop_remainder, use_legacy_function=False):
+    self._input_dataset = input_dataset
+
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func,
+        "tf.data.experimental.map_and_batch()",
+        dataset=input_dataset,
+        use_legacy_function=use_legacy_function)
+    self._batch_size_t = ops.convert_to_tensor(
+        batch_size, dtype=dtypes.int64, name="batch_size")
+    self._num_parallel_calls_t = ops.convert_to_tensor(
+        num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
+    self._drop_remainder_t = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+
+    constant_drop_remainder = tensor_util.constant_value(self._drop_remainder_t)
+    # pylint: disable=protected-access
+    if constant_drop_remainder:
+      # NOTE(mrry): `constant_drop_remainder` may be `None` (unknown statically)
+      # or `False` (explicitly retaining the remainder).
+      # pylint: disable=g-long-lambda
+      self._element_spec = nest.map_structure(
+          lambda component_spec: component_spec._batch(
+              tensor_util.constant_value(self._batch_size_t)),
+          self._map_func.output_structure)
+    else:
+      self._element_spec = nest.map_structure(
+          lambda component_spec: component_spec._batch(None),
+          self._map_func.output_structure)
+    # pylint: enable=protected-access
+    variant_tensor = ged_ops.map_and_batch_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        f=self._map_func.function,
+        batch_size=self._batch_size_t,
+        num_parallel_calls=self._num_parallel_calls_t,
+        drop_remainder=self._drop_remainder_t,
+        preserve_cardinality=True,
+        **self._flat_structure)
+    super(_MapAndBatchDataset, self).__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._element_spec

tensorflow/python/data/experimental/ops/cardinality.py

@@ -0,0 +1,113 @@
+# 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.
+# ==============================================================================
+"""Cardinality analysis of `Dataset` objects."""
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+INFINITE = -1
+UNKNOWN = -2
+tf_export("data.experimental.INFINITE_CARDINALITY").export_constant(
+    __name__, "INFINITE")
+tf_export("data.experimental.UNKNOWN_CARDINALITY").export_constant(
+    __name__, "UNKNOWN")
+
+
+# TODO(b/157691652): Deprecate this method after migrating users to the new API.
+@tf_export("data.experimental.cardinality")
+def cardinality(dataset):
+  """Returns the cardinality of `dataset`, if known.
+
+  The operation returns the cardinality of `dataset`. The operation may return
+  `tf.data.experimental.INFINITE_CARDINALITY` if `dataset` contains an infinite
+  number of elements or `tf.data.experimental.UNKNOWN_CARDINALITY` if the
+  analysis fails to determine the number of elements in `dataset` (e.g. when the
+  dataset source is a file).
+
+  >>> dataset = tf.data.Dataset.range(42)
+  >>> print(tf.data.experimental.cardinality(dataset).numpy())
+  42
+  >>> dataset = dataset.repeat()
+  >>> cardinality = tf.data.experimental.cardinality(dataset)
+  >>> print((cardinality == tf.data.experimental.INFINITE_CARDINALITY).numpy())
+  True
+  >>> dataset = dataset.filter(lambda x: True)
+  >>> cardinality = tf.data.experimental.cardinality(dataset)
+  >>> print((cardinality == tf.data.experimental.UNKNOWN_CARDINALITY).numpy())
+  True
+
+  Args:
+    dataset: A `tf.data.Dataset` for which to determine cardinality.
+
+  Returns:
+    A scalar `tf.int64` `Tensor` representing the cardinality of `dataset`. If
+    the cardinality is infinite or unknown, the operation returns the named
+    constant `INFINITE_CARDINALITY` and `UNKNOWN_CARDINALITY` respectively.
+  """
+
+  return gen_dataset_ops.dataset_cardinality(dataset._variant_tensor)  # pylint: disable=protected-access
+
+
+@tf_export("data.experimental.assert_cardinality")
+def assert_cardinality(expected_cardinality):
+  """Asserts the cardinality of the input dataset.
+
+  NOTE: The following assumes that "examples.tfrecord" contains 42 records.
+
+  >>> dataset = tf.data.TFRecordDataset("examples.tfrecord")
+  >>> cardinality = tf.data.experimental.cardinality(dataset)
+  >>> print((cardinality == tf.data.experimental.UNKNOWN_CARDINALITY).numpy())
+  True
+  >>> dataset = dataset.apply(tf.data.experimental.assert_cardinality(42))
+  >>> print(tf.data.experimental.cardinality(dataset).numpy())
+  42
+
+  Args:
+    expected_cardinality: The expected cardinality of the input dataset.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+
+  Raises:
+    FailedPreconditionError: The assertion is checked at runtime (when iterating
+      the dataset) and an error is raised if the actual and expected cardinality
+      differ.
+  """
+  def _apply_fn(dataset):
+    return _AssertCardinalityDataset(dataset, expected_cardinality)
+
+  return _apply_fn
+
+
+class _AssertCardinalityDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that assert the cardinality of its input."""
+
+  def __init__(self, input_dataset, expected_cardinality):
+    self._input_dataset = input_dataset
+    self._expected_cardinality = ops.convert_to_tensor(
+        expected_cardinality, dtype=dtypes.int64, name="expected_cardinality")
+
+    # pylint: enable=protected-access
+    variant_tensor = ged_ops.assert_cardinality_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._expected_cardinality,
+        **self._flat_structure)
+    super(_AssertCardinalityDataset, self).__init__(input_dataset,
+                                                    variant_tensor)

tensorflow/python/data/experimental/ops/compression_ops.py

@@ -0,0 +1,51 @@
+# 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.
+# ==============================================================================
+"""Ops for compressing and uncompressing dataset elements."""
+from tensorflow.python.data.util import structure
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def compress(element):
+  """Compress a dataset element.
+
+  Args:
+    element: A nested structure of types supported by Tensorflow.
+
+  Returns:
+    A variant tensor representing the compressed element. This variant can be
+    passed to `uncompress` to get back the original element.
+  """
+  element_spec = structure.type_spec_from_value(element)
+  tensor_list = structure.to_tensor_list(element_spec, element)
+  return ged_ops.compress_element(tensor_list)
+
+
+def uncompress(element, output_spec):
+  """Uncompress a compressed dataset element.
+
+  Args:
+    element: A scalar variant tensor to uncompress. The element should have been
+      created by calling `compress`.
+    output_spec: A nested structure of `tf.TypeSpec` representing the type(s) of
+      the uncompressed element.
+
+  Returns:
+    The uncompressed element.
+  """
+  flat_types = structure.get_flat_tensor_types(output_spec)
+  flat_shapes = structure.get_flat_tensor_shapes(output_spec)
+  tensor_list = ged_ops.uncompress_element(
+      element, output_types=flat_types, output_shapes=flat_shapes)
+  return structure.from_tensor_list(output_spec, tensor_list)

tensorflow/python/data/experimental/ops/counter.py

@@ -0,0 +1,84 @@
+# 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.
+# ==============================================================================
+"""The Counter Dataset."""
+from tensorflow.python import tf2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.Counter", v1=[])
+@deprecation.deprecated(None, "Use `tf.data.Dataset.counter(...)` instead.")
+def CounterV2(start=0, step=1, dtype=dtypes.int64):
+  """Creates a `Dataset` that counts from `start` in steps of size `step`.
+
+  Unlike `tf.data.Dataset.range` which will stop at some ending number,
+  `Counter` will produce elements indefinitely.
+
+  >>> dataset = tf.data.experimental.Counter().take(5)
+  >>> [a.item() for a in dataset.as_numpy_iterator()]
+  [0, 1, 2, 3, 4]
+  >>> dataset.element_spec
+  TensorSpec(shape=(), dtype=tf.int64, name=None)
+  >>> dataset = tf.data.experimental.Counter(dtype=tf.int32)
+  >>> dataset.element_spec
+  TensorSpec(shape=(), dtype=tf.int32, name=None)
+  >>> dataset = tf.data.experimental.Counter(start=2).take(5)
+  >>> [a.item() for a in dataset.as_numpy_iterator()]
+  [2, 3, 4, 5, 6]
+  >>> dataset = tf.data.experimental.Counter(start=2, step=5).take(5)
+  >>> [a.item() for a in dataset.as_numpy_iterator()]
+  [2, 7, 12, 17, 22]
+  >>> dataset = tf.data.experimental.Counter(start=10, step=-1).take(5)
+  >>> [a.item() for a in dataset.as_numpy_iterator()]
+  [10, 9, 8, 7, 6]
+
+  Args:
+    start: (Optional.) The starting value for the counter. Defaults to 0.
+    step: (Optional.) The step size for the counter. Defaults to 1.
+    dtype: (Optional.) The data type for counter elements. Defaults to
+      `tf.int64`.
+
+  Returns:
+    A `Dataset` of scalar `dtype` elements.
+  """
+  return dataset_ops.Dataset.counter(start, step, dtype)
+
+
+@tf_export(v1=["data.experimental.Counter"])
+@deprecation.deprecated(None, "Use `tf.data.Dataset.counter(...)` instead.")
+def CounterV1(start=0, step=1, dtype=dtypes.int64):
+  return dataset_ops.DatasetV1Adapter(CounterV2(start, step, dtype))
+
+
+CounterV1.__doc__ = CounterV2.__doc__
+
+if tf2.enabled():
+  Counter = CounterV2
+else:
+  Counter = CounterV1
+
+
+def _tf2_callback():  # pylint: disable=invalid-name
+  global Counter
+  if tf2.enabled():
+    Counter = CounterV2
+  else:
+    Counter = CounterV1
+
+
+v2_compat.register_data_v2_callback(_tf2_callback)

tensorflow/python/data/experimental/ops/data_service_ops_test.py

@@ -0,0 +1,99 @@
+# 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 HashableElementSpec and related functionality."""
+
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+class HashableElementSpecTest(test.TestCase):
+
+  def testEqual(self):
+    spec1 = data_service_ops.HashableElementSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    )
+    spec2 = data_service_ops.HashableElementSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    )
+    self.assertEqual(spec1, spec2)
+    self.assertEqual(hash(spec1), hash(spec2))
+
+  def testNotEqual(self):
+    spec1 = data_service_ops.HashableElementSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    )
+    spec2 = data_service_ops.HashableElementSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64)
+    )
+    self.assertNotEqual(spec1, spec2)
+    self.assertNotEqual(hash(spec1), hash(spec2))
+
+  def testNotEqualOtherType(self):
+    spec1 = data_service_ops.HashableElementSpec(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    )
+    self.assertNotEqual(spec1, 123)
+
+  def testGetUncompressFuncCache(self):
+    spec1 = tensor_spec.TensorSpec(shape=(), dtype=dtypes.int32)
+    spec2 = tensor_spec.TensorSpec(shape=(), dtype=dtypes.int64)
+    data_service_ops._get_uncompress_func.cache_clear()
+    func1 = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec1)
+    )
+    func2 = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec1)
+    )
+    func3 = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec2)
+    )
+    self.assertIs(func1, func2)
+    self.assertIsNot(func1, func3)
+
+  def testRaggedTensorUncompressFuncCache(self):
+    spec1 = ragged_tensor.RaggedTensorSpec(
+        [3, None], dtypes.int32, 1, dtypes.int64
+    )
+    spec1_long_format = ragged_tensor.RaggedTensorSpec(
+        tensor_shape.TensorShape(
+            [tensor_shape.Dimension(3), tensor_shape.Dimension(None)]
+        ),
+        dtypes.int32,
+        1,
+        dtypes.int64,
+    )
+    spec2 = ragged_tensor.RaggedTensorSpec(
+        ([3, None]), dtypes.int64, 1, dtypes.int64
+    )
+    data_service_ops._get_uncompress_func.cache_clear()
+    func1a = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec1)
+    )
+    func1b = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec1_long_format)
+    )
+    func2 = data_service_ops._get_uncompress_func(
+        data_service_ops.HashableElementSpec(spec2)
+    )
+    self.assertIs(func1a, func1b)
+    self.assertIsNot(func1a, func2)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/experimental/ops/distribute.py

@@ -0,0 +1,399 @@
+# 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.
+# ==============================================================================
+"""Distribution Strategy-related dataset transformations."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops.options import ExternalStatePolicy
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.types import data as data_types
+from tensorflow.python.util.tf_export import tf_export
+
+SHARD_HINT = -1
+tf_export("data.experimental.SHARD_HINT").export_constant(
+    __name__, "SHARD_HINT")
+
+
+class _AutoShardDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that shards the `Dataset` automatically.
+
+  This dataset takes in an existing dataset and tries to automatically figure
+  out how to shard the dataset in a multi-worker scenario using graph rewrites.
+
+  If the AutoShardPolicy is set to FILE, it walks up the dataset graph until
+  it finds a reader dataset, then inserts a ShardDataset op before that node
+  so that each worker only sees some files.
+
+  If the AutoShardPolicy is set to DATA, it inserts a ShardDataset op at the
+  end of the input pipeline, before any terminal PrefetchDataset if there is
+  one. Additionally, if there is a RebatchDatasetV2 in the input pipeline, it
+  is written to legacy RebatchDataset for correctness reasons, since
+  RebatchDatasetV2 is incompatible with data sharding.
+
+  If the AutoShardPolicy is set to AUTO, it tries to do file-based sharding.
+  If it cannot find a reader dataset, it falls back to doing data-based
+  sharding.
+
+  If the AutoShardPolicy is set to OFF, it does nothing.
+
+  Attributes:
+    num_workers: Total number of workers to shard this dataset across.
+    index: The current worker index (out of the total number of workers) this
+      dataset is for.
+    num_replicas: The total number of replicas across all workers. This is used
+      only when sharding by data (either DATA or AUTO) in order to rewrite
+      RebatchDatasetV2 to RebatchDataset.
+
+  Raises:
+    NotFoundError: If we cannot find a suitable reader dataset to begin
+      automatically sharding the dataset.
+  """
+
+  def __init__(self, input_dataset, num_workers, index, num_replicas=None):
+    self._input_dataset = input_dataset
+
+    self._element_spec = input_dataset.element_spec
+    variant_tensor = ged_ops.auto_shard_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        num_workers=num_workers,
+        index=index,
+        auto_shard_policy=int(
+            input_dataset.options().experimental_distribute.auto_shard_policy),
+        num_replicas=num_replicas,
+        **self._flat_structure)
+    super(_AutoShardDataset, self).__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+def _AutoShardDatasetV1(input_dataset, num_workers, index, num_replicas=None):  # pylint: disable=invalid-name
+  return dataset_ops.DatasetV1Adapter(
+      _AutoShardDataset(input_dataset, num_workers, index, num_replicas))
+
+
+class _LegacyRebatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that divides its input batches into `num_replicas` sub-batches.
+
+  For each batch in the input dataset, _LegacyRebatchDataset will produce
+  `num_replicas` smaller batches whose sizes add up to the original batch size.
+
+  For example:
+
+  ```python
+  ds = tf.data.Dataset.range(8)
+  ds = ds.batch(4)
+  ds = _LegacyRebatchDataset(ds, num_replicas=3)
+  for elem in ds:
+    print(elem)
+  >> [0, 1], [2, 3], [], [4, 5], [6, 7], []
+  ```
+  """
+
+  def __init__(self, input_dataset, num_replicas):
+    """Creates a _LegacyRebatchDataset.
+
+    Args:
+      input_dataset: `Dataset` to rebatch.
+      num_replicas: A `tf.int64` scalar, representing the number of sub-batches
+        to split each batch from `input_dataset` into.
+    """
+
+    def recalculate_batch_size(type_spec):
+      """Recalculates the output_shape after dividing it by num_replicas."""
+      output_shape = type_spec._to_legacy_output_shapes()  # pylint: disable=protected-access
+      if not isinstance(output_shape, tensor_shape.TensorShape):
+        return None
+
+      # If the output shape is unknown, we set the batch dimension to unknown.
+      if output_shape.rank is None:
+        return None
+
+      if len(output_shape) < 1:
+        raise ValueError(
+            "Invalid `input_dataset`. Expected a dataset whose elements "
+            "have rank >= 1 but found a dataset whose elements are scalars. "
+            "Fix the issue by adding the `batch` transformation to the "
+            "dataset.")
+      output_dims = [d.value for d in output_shape.dims]
+
+      if output_dims[0] is not None and output_dims[0] % num_replicas == 0:
+        return output_dims[0] // num_replicas
+
+      # Set the batch dimension to unknown. If the global batch size does not
+      # divide num_replicas evenly, the minibatches may have different sizes.
+      return None
+
+    def rebatch(type_spec):
+      # pylint: disable=protected-access
+      batch_size = recalculate_batch_size(type_spec)
+      return type_spec._unbatch()._batch(batch_size)
+      # pylint: enable=protected-access
+
+    self._element_spec = nest.map_structure(
+        rebatch, dataset_ops.get_structure(input_dataset))
+
+    # auto_shard rewrite assumes that there's normalize_to_dense before
+    # rebatch_dataset.
+    # LINT.IfChange
+    input_dataset = dataset_ops.normalize_to_dense(input_dataset)
+    variant_tensor = ged_ops.rebatch_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        num_replicas=num_replicas,
+        **self._flat_structure)
+    # LINT.ThenChange(//tensorflow/core/grappler/optimizers/data/auto_shard.cc)
+    super(_LegacyRebatchDataset, self).__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+class _RemoteDataset(dataset_ops.DatasetSource):
+  """Creates a dataset on a given `device` given a graph def."""
+
+  def __init__(self, graph_def, device, element_spec):
+    self._elem_spec = element_spec
+    with ops.device(device):
+      variant_tensor = ged_ops.dataset_from_graph(graph_def)
+    super(_RemoteDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._elem_spec
+
+
+def replicate(dataset, devices):
+  """A transformation that replicates `dataset` onto a list of devices.
+
+  Args:
+    dataset: A `tf.data.Dataset` object.
+    devices: A list of devices to replicate the dataset on.
+
+  Returns:
+    A dictionary mapping device name to a dataset on that device.
+  """
+  if not isinstance(dataset, data_types.DatasetV2):
+    raise TypeError(
+        f"Invalid `dataset`. Expected a `tf.data.Dataset` object but "
+        f"got {type(dataset)}.")
+
+  # pylint: disable=protected-access
+  dataset_device = dataset._variant_tensor.device
+
+  datasets = {}
+  if len(devices) == 1 and devices[0] == dataset_device:
+    datasets[devices[0]] = dataset
+    return datasets
+
+  with ops.colocate_with(dataset._variant_tensor):
+    dataset = dataset._apply_debug_options()
+    graph_def = dataset._as_serialized_graph(
+        strip_device_assignment=True,
+        external_state_policy=ExternalStatePolicy.WARN)
+  for device in devices:
+    ds = _RemoteDataset(graph_def, device, dataset.element_spec)
+    datasets[device] = ds
+  return datasets
+
+
+def batch_sizes_for_worker(global_batch_size, num_workers,
+                           num_replicas_per_worker, worker_index):
+  """Determines how to rebatch a dataset for the given worker.
+
+  Given the global batch size, number of workers, number of replicas per worker,
+  and worker index, returns the correct batch sizes for rebatching a dataset
+  on worker `worker_index` of `num_workers`, such that each global step (across
+  all workers and replicas) will consume global_batch_size elements. The
+  returned value should be passed as the `batch_sizes` input parameter to
+  `tf.data.experimental.rebatch()`. The returned batch sizes meet the following
+  constraints:
+
+  Let G = global_batch_size, W = num_workers, R = num_replicas_per_worker
+  (A) for any worker, len(batch_sizes) = W * R
+  (B) for any worker, sum(batch_sizes) == G
+  (C) for any global step (i.e. R iterations on each worker), the sum of batches
+      consumed by replicas across all workers is G.
+  (D) any two batch sizes of any two replicas differs by at most one.
+
+  For example, suppose we have G = 7, W = 2, R = 2, and suppose we have two
+  files which each contain 7 elements:
+
+  ```python
+  # WORKER 0
+  batch_sizes_0 = batch_sizes_for_worker(global_batch_size=global_batch_size,
+                                         num_workers=2,
+                                         num_replicas_per_worker=2,
+                                         worker_index=0)
+  print(batch_sizes_0)
+  >> [2, 2, 2, 1]
+
+  dataset_0 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
+  dataset_0 = dataset_0.shard(num_shards, index=0)
+  dataset_0 = dataset_0.batch(7)
+  dataset_0 = dataset_0.apply(tf.data.experimental.rebatch(batch_sizes_0))
+  for elem in dataset_0:
+    print(elem)
+  >> [[A0, A1], [A2, A3], [A4, A5], [A6]]
+
+  # WORKER 1
+  batch_sizes_1 = batch_sizes_for_worker(global_batch_size=global_batch_size,
+                                         num_workers=2,
+                                         num_replicas_per_worker=2,
+                                         worker_index=1)
+  print(batch_sizes_1)
+  >> [2, 1, 2, 2]
+
+  dataset_1 = tf.data.Dataset.from_tensor_slices(["file_a", "file_b"])
+  dataset_1 = dataset_1.shard(num_shards, index=1)
+  dataset_1 = dataset_1.batch(7)
+  dataset_1 = dataset_1.apply(tf.data.experimental.rebatch(batch_sizes_1))
+  for elem in dataset_1:
+    print(elem)
+  >> [[B0, B1], [B2], [B3, B4], [B5, B6]]
+  ```
+
+  The above example will produce the following elements:
+
+  Step 1:
+    Worker 0 Replica 0: [A0, A1]
+    Worker 0 Replica 1: [A2, A3]
+    Worker 1 Replica 0: [B0, B1]
+    Worker 1 Replica 1: [B2]
+  Total batch size = 7
+
+  Step 2:
+    Worker 0 Replica 0: [A4, A5]
+    Worker 0 Replica 1: [A6]
+    Worker 1 Replica 0: [B3, B4]
+    Worker 1 Replica 1: [B5, B6]
+  Total batch size = 7
+
+  Args:
+    global_batch_size: A `tf.int64` scalar, representing the global batch size.
+    num_workers: An integer representing the number of workers the dataset will
+      be distributed across.
+    num_replicas_per_worker: An integer representing the number of replicas per
+      worker. All workers are assumed to have the same number of replicas.
+    worker_index: An integer index of the worker to be rebatched.
+
+  Returns:
+    A `tf.int64` vector, representing the batch sizes to rebatch the dataset
+    into.
+  """
+  # Constraint (A)
+  num_subbatches = num_workers * num_replicas_per_worker
+
+  offset = worker_index * num_replicas_per_worker
+
+  const_value = tensor_util.constant_value(global_batch_size)
+  if const_value is not None:
+    # Use the constant global batch size for further calculations
+    global_batch_size = const_value
+
+  # Let N = W * R. Constraint (B) and (D) jointly mean that the iterations
+  # should have batch size either floor(B/N) or ceil(B/N). Namely, of the N
+  # subbatches a batch is split into, B - N * floor(B/N) of them will have size
+  # ceil(B/N), and the rest will have size floor(B/N).
+  floor = global_batch_size // num_subbatches
+  num_ceil = global_batch_size - (num_subbatches * floor)
+
+  # For worker 0, we assign the first num_ceil subbatches to have size
+  # ceil(B/N), and the remainder to have size floor(B/N). The other workers will
+  # each be offset by R * worker_index in order to meet constraint (C).
+  if const_value is not None:
+    # If the global batch size is a known constant value, we return a constant
+    # tensor directly instead of manipulating it with TF ops. This allows for
+    # better downstream shape inference.
+    worker_0 = [floor + 1] * num_ceil + [floor] * (num_subbatches - num_ceil)
+    return ops.convert_to_tensor(
+        worker_0[offset:] + worker_0[:offset],
+        dtype=dtypes.int64,
+        name="batch_sizes")
+
+  worker_0 = array_ops.ones(num_subbatches, dtype=dtypes.int64)
+  worker_0 = floor * worker_0 + array_ops.concat([
+      array_ops.ones(num_ceil, dtype=dtypes.int64),
+      array_ops.zeros(num_subbatches - num_ceil, dtype=dtypes.int64)
+  ],
+                                                 axis=0)
+
+  return array_ops.concat([worker_0[offset:], worker_0[:offset]], axis=0)
+
+
+def compute_batch_size(dataset):
+  """An operation that returns the batch size of the dataset.
+
+  This op tries to infer the batch size statically by walking up the dataset
+  tree from the final dataset node and returning the batch size of the first
+  batching dataset (such as from .batch() and .padded_batch()) that it
+  encounters. This differs from using the `element_spec` of a dataset in that it
+  does not account for partial batches.
+
+  This operation may fail if it encounters contradictory batch sizes (for
+  example, if the dataset is created by zipping together two datasets with
+  different batch sizes), if there are no explicit batching transformations, or
+  if there are operations downstream from the batching transformation that may
+  modify its batch size. In these cases, it returns a -1.
+
+  Args:
+    dataset: A `tf.data.Dataset` object.
+
+  Returns:
+    A `tf.int64` Tensor representing the batch size of the dataset sans partial
+    batches. If this cannot be inferred statically, the value of this tensor
+    will be -1.
+  """
+
+  def get_static_batch_dim(type_spec):
+    try:
+      output_shape = type_spec._to_legacy_output_shapes()  # pylint: disable=protected-access
+    except NotImplementedError:
+      return None
+    if not isinstance(output_shape, tensor_shape.TensorShape):
+      return None
+    if output_shape.rank is None:
+      return None
+    return output_shape.dims[0].value
+
+  batch_dims = [
+      get_static_batch_dim(type_spec)
+      for type_spec in nest.flatten(dataset_ops.get_structure(dataset))
+  ]
+
+  if all(d is not None for d in batch_dims):
+
+    if all(d == batch_dims[0] for d in batch_dims):
+      # If all batch dimensions are known and equal, return that directly.
+      batch_dim = batch_dims[0]
+    else:
+      # If all batch dimensions are known but not all equal, return -1.
+      batch_dim = -1
+
+    return constant_op.constant(
+        batch_dim, dtype=dtypes.int64, name="static_batch_size")
+
+  # If any batch dimensions are unknown, use compute_batch_size op.
+  return ged_ops.compute_batch_size(dataset._variant_tensor)  # pylint: disable=protected-access
+
+
+_AutoShardDatasetV1.__doc__ = _AutoShardDataset.__doc__

tensorflow/python/data/experimental/ops/distributed_save_op.py

@@ -0,0 +1,111 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Distributed saving of a dataset to disk."""
+
+from typing import Optional
+
+from tensorflow.core.protobuf import snapshot_pb2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+# TODO(b/238903802): Use TypeSpec serialization methods directly.
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.distributed_save")
+def distributed_save(
+    dataset: dataset_ops.Dataset,
+    path: str,
+    data_service_address: str,
+    compression: str = "AUTO",
+) -> Optional[ops.OperationType]:
+  """Initiates the process of saving a dataset to disk using tf.data service.
+
+  The op uses tf.data service
+  (https://www.tensorflow.org/api_docs/python/tf/data/experimental/service) to
+  write a dataset snapshot. Returns immediately after submitting the request.
+  Does not wait for the snapshot to be finished. Requires that the tf.data
+  service run a fixed number of worker replicas.
+
+  To load the snapshot, users may optionally pass `wait=True` to
+  `tf.data.Dataset.load` so it can read snapshots as they are being written.
+
+  Example usage:
+
+  >>> import os
+  >>> import tempfile
+
+  >>> # Runs tf.data service.
+  >>> tempdir = tempfile.gettempdir()
+  >>> dispatcher = tf.data.experimental.service.DispatchServer(
+  ...     tf.data.experimental.service.DispatcherConfig(
+  ...         fault_tolerant_mode=True,
+  ...         work_dir=os.path.join(tempdir, "work_dir")))
+  >>> dispatcher_address = dispatcher.target.split("://")[1]
+  >>> worker = tf.data.experimental.service.WorkerServer(
+  ...     tf.data.experimental.service.WorkerConfig(
+  ...         dispatcher_address=dispatcher_address))
+
+  >>> # Writes dataset snapshot.
+  >>> path = os.path.join(tempdir, "dataset_snapshot")
+  >>> dataset = tf.data.Dataset.range(1)
+  >>> tf.data.experimental.distributed_save(dataset, path, dispatcher_address)
+
+  >>> # Loads a dataset snapshot.
+  >>> loaded_dataset = tf.data.Dataset.load(path, wait=True)
+  >>> for elem in loaded_dataset:
+  ...   print(elem)
+  tf.Tensor(0, shape=(), dtype=int64)
+
+  Args:
+    dataset: The `tf.data.Dataset` to save.
+    path: The directory path to save the dataset. Requires that:
+      - The directory does not exist and will create the directory.
+      - The file system supports atomic move (rename).
+    data_service_address: tf.data service dispatcher address.
+    compression: (Optional.) Whether and how to compress the `dataset` snapshot.
+      If `"AUTO"`, the tf.data runtime decides which algorithm to use. If
+      `"GZIP"` or `"SNAPPY"`, that specific algorithm is used.  If `None`, the
+      `dataset` snapshot is not compressed.
+
+  Returns:
+    An operation which when executed performs the distributed save.
+
+  Raises:
+    ValueError: If `dispatcher_address` is invalid.
+    tf.errors.AlreadyExistsError: If the snapshot has already started or has
+      finished.
+    tf.errors.FailedPreconditionError: If the file system does not support
+      atomic move (rename).
+    tf.errors.InvalidArgumentError: If tf.data service is not running in the
+      fault tolerant mode.
+  """
+  if not isinstance(data_service_address, str):
+    raise ValueError("`data_service_address` must be a string, but is a "
+                     f"{type(data_service_address)} ({data_service_address}")
+  if not data_service_address:
+    raise ValueError("`data_service_address` must not be empty")
+
+  metadata = snapshot_pb2.DistributedSnapshotMetadata(
+      element_spec=nested_structure_coder.encode_structure(
+          dataset.element_spec).SerializeToString(),
+      compression=compression)
+
+  return gen_experimental_dataset_ops.distributed_save(
+      dataset._variant_tensor,  # pylint: disable=protected-access
+      directory=path,
+      address=data_service_address,
+      metadata=metadata.SerializeToString())

tensorflow/python/data/experimental/ops/enumerate_ops.py

@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+"""Enumerate dataset transformations."""
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated(None, "Use `tf.data.Dataset.enumerate()`.")
+@tf_export("data.experimental.enumerate_dataset")
+def enumerate_dataset(start=0):
+  """A transformation that enumerates the elements of a dataset.
+
+  It is similar to python's `enumerate`.
+  For example:
+
+  ```python
+  # NOTE: The following examples use `{ ... }` to represent the
+  # contents of a dataset.
+  a = { 1, 2, 3 }
+  b = { (7, 8), (9, 10) }
+
+  # The nested structure of the `datasets` argument determines the
+  # structure of elements in the resulting dataset.
+  a.apply(tf.data.experimental.enumerate_dataset(start=5))
+  => { (5, 1), (6, 2), (7, 3) }
+  b.apply(tf.data.experimental.enumerate_dataset())
+  => { (0, (7, 8)), (1, (9, 10)) }
+  ```
+
+  Args:
+    start: A `tf.int64` scalar `tf.Tensor`, representing the start value for
+      enumeration.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+  """
+
+  def _apply_fn(dataset):
+    return dataset.enumerate(start)
+
+  return _apply_fn

tensorflow/python/data/experimental/ops/error_ops.py

@@ -0,0 +1,51 @@
+# 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.
+# ==============================================================================
+"""Ignore_errors dataset transformations."""
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.ignore_errors")
+@deprecation.deprecated(None, "Use `tf.data.Dataset.ignore_errors` instead.")
+def ignore_errors(log_warning=False):
+  """Creates a `Dataset` from another `Dataset` and silently ignores any errors.
+
+  Use this transformation to produce a dataset that contains the same elements
+  as the input, but silently drops any elements that caused an error. For
+  example:
+
+  ```python
+  dataset = tf.data.Dataset.from_tensor_slices([1., 2., 0., 4.])
+
+  # Computing `tf.debugging.check_numerics(1. / 0.)` will raise an
+  InvalidArgumentError.
+  dataset = dataset.map(lambda x: tf.debugging.check_numerics(1. / x, "error"))
+
+  # Using `ignore_errors()` will drop the element that causes an error.
+  dataset =
+      dataset.apply(tf.data.experimental.ignore_errors())  # ==> {1., 0.5, 0.2}
+  ```
+  Args:
+     log_warning: (Optional.) A 'tf.bool' scalar indicating whether ignored
+      errors should be logged to stderr. Defaults to 'False'.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    `tf.data.Dataset.apply`.
+  """
+  def _apply_fn(dataset):
+    return dataset.ignore_errors(log_warning)
+
+  return _apply_fn

tensorflow/python/data/experimental/ops/from_list.py

@@ -0,0 +1,122 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Python API for creating a dataset from a list."""
+
+import itertools
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+class _ListDataset(dataset_ops.DatasetSource):
+  """A `Dataset` of elements from a list."""
+
+  def __init__(self, elements, name=None):
+    if not elements:
+      raise ValueError(
+          "Invalid `elements`. `elements` should not be empty. If you want an"
+          " empty dataset, use `tf.data.Dataset.range(0)`."
+      )
+    if not isinstance(elements, list):
+      raise ValueError("Invalid `elements`. `elements` must be a list.")
+
+    elements = [structure.normalize_element(element) for element in elements]
+    type_specs = [
+        structure.type_spec_from_value(element) for element in elements
+    ]
+
+    # Check that elements have same nested structure.
+    num_elements = len(elements)
+    for i in range(1, num_elements):
+      nest.assert_same_structure(type_specs[0], type_specs[i])
+
+    # Infer elements' supershape.
+    flattened_type_specs = [nest.flatten(type_spec) for type_spec in type_specs]
+    num_tensors_per_element = len(flattened_type_specs[0])
+    flattened_structure = [None] * num_tensors_per_element
+    for i in range(num_tensors_per_element):
+      flattened_structure[i] = flattened_type_specs[0][i]
+      for j in range(1, num_elements):
+        flattened_structure[i] = flattened_structure[
+            i].most_specific_common_supertype([flattened_type_specs[j][i]])
+
+    if not isinstance(type_specs[0], dataset_ops.DatasetSpec):
+      self._tensors = list(
+          itertools.chain.from_iterable(
+              [nest.flatten(element) for element in elements]))
+    else:
+      self._tensors = [x._variant_tensor for x in elements]
+    self._structure = nest.pack_sequence_as(type_specs[0], flattened_structure)
+    self._name = name
+    variant_tensor = gen_experimental_dataset_ops.list_dataset(
+        self._tensors,
+        output_types=self._flat_types,
+        output_shapes=self._flat_shapes,
+        metadata=self._metadata.SerializeToString())
+    super(_ListDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+
+@tf_export("data.experimental.from_list")
+def from_list(elements, name=None):
+  """Creates a `Dataset` comprising the given list of elements.
+
+  The returned dataset will produce the items in the list one by one. The
+  functionality is identical to `Dataset.from_tensor_slices` when elements are
+  scalars, but different when elements have structure. Consider the following
+  example.
+
+  >>> dataset = tf.data.experimental.from_list([(1, 'a'), (2, 'b'), (3, 'c')])
+  >>> [(n.item(), s) for n, s in dataset.as_numpy_iterator()]
+  [(1, b'a'), (2, b'b'), (3, b'c')]
+
+  To get the same output with `from_tensor_slices`, the data needs to be
+  reorganized:
+
+  >>> dataset = tf.data.Dataset.from_tensor_slices(([1, 2, 3], ['a', 'b', 'c']))
+  >>> [(n.item(), s) for n, s in dataset.as_numpy_iterator()]
+  [(1, b'a'), (2, b'b'), (3, b'c')]
+
+  Unlike `from_tensor_slices`, `from_list` supports non-rectangular input:
+
+  >>> dataset = tf.data.experimental.from_list([[1], [2, 3]])
+  >>> list(dataset.as_numpy_iterator())
+  [array([1], dtype=int32), array([2, 3], dtype=int32)]
+
+  Achieving the same with `from_tensor_slices` requires the use of ragged
+  tensors.
+
+  `from_list` can be more performant than `from_tensor_slices` in some cases,
+  since it avoids the need for data slicing each epoch. However, it can also be
+  less performant, because data is stored as many small tensors rather than a
+  few large tensors as in `from_tensor_slices`. The general guidance is to
+  prefer `from_list` from a performance perspective when the number of elements
+  is small (less than 1000).
+
+  Args:
+    elements: A list of elements whose components have the same nested
+      structure.
+    name: (Optional.) A name for the tf.data operation.
+
+  Returns:
+    Dataset: A `Dataset` of the `elements`.
+  """
+  return _ListDataset(elements, name)

tensorflow/python/data/experimental/ops/get_single_element.py

@@ -0,0 +1,109 @@
+# 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.
+# ==============================================================================
+"""Python wrappers for Datasets and Iterators."""
+from tensorflow.python.types import data as data_types
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@deprecation.deprecated(None, "Use `tf.data.Dataset.get_single_element()`.")
+@tf_export("data.experimental.get_single_element")
+def get_single_element(dataset):
+  """Returns the single element of the `dataset` as a nested structure of tensors.
+
+  The function enables you to use a `tf.data.Dataset` in a stateless
+  "tensor-in tensor-out" expression, without creating an iterator.
+  This facilitates the ease of data transformation on tensors using the
+  optimized `tf.data.Dataset` abstraction on top of them.
+
+  For example, lets consider a `preprocessing_fn` which would take as an
+  input the raw features and returns the processed feature along with
+  it's label.
+
+  ```python
+  def preprocessing_fn(raw_feature):
+    # ... the raw_feature is preprocessed as per the use-case
+    return feature
+
+  raw_features = ...  # input batch of BATCH_SIZE elements.
+  dataset = (tf.data.Dataset.from_tensor_slices(raw_features)
+             .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
+             .batch(BATCH_SIZE))
+
+  processed_features = tf.data.experimental.get_single_element(dataset)
+  ```
+
+  In the above example, the `raw_features` tensor of length=BATCH_SIZE
+  was converted to a `tf.data.Dataset`. Next, each of the `raw_feature` was
+  mapped using the `preprocessing_fn` and the processed features were
+  grouped into a single batch. The final `dataset` contains only one element
+  which is a batch of all the processed features.
+
+  NOTE: The `dataset` should contain only one element.
+
+  Now, instead of creating an iterator for the `dataset` and retrieving the
+  batch of features, the `tf.data.experimental.get_single_element()` function
+  is used to skip the iterator creation process and directly output the batch
+  of features.
+
+  This can be particularly useful when your tensor transformations are
+  expressed as `tf.data.Dataset` operations, and you want to use those
+  transformations while serving your model.
+
+  # Keras
+
+  ```python
+
+  model = ... # A pre-built or custom model
+
+  class PreprocessingModel(tf.keras.Model):
+    def __init__(self, model):
+      super().__init__(self)
+      self.model = model
+
+    @tf.function(input_signature=[...])
+    def serving_fn(self, data):
+      ds = tf.data.Dataset.from_tensor_slices(data)
+      ds = ds.map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
+      ds = ds.batch(batch_size=BATCH_SIZE)
+      return tf.argmax(
+        self.model(tf.data.experimental.get_single_element(ds)),
+        axis=-1
+      )
+
+  preprocessing_model = PreprocessingModel(model)
+  your_exported_model_dir = ... # save the model to this path.
+  tf.saved_model.save(preprocessing_model, your_exported_model_dir,
+                signatures={'serving_default': preprocessing_model.serving_fn})
+  ```
+
+  Args:
+    dataset: A `tf.data.Dataset` object containing a single element.
+
+  Returns:
+    A nested structure of `tf.Tensor` objects, corresponding to the single
+    element of `dataset`.
+
+  Raises:
+    TypeError: if `dataset` is not a `tf.data.Dataset` object.
+    InvalidArgumentError: (at runtime) if `dataset` does not contain exactly
+      one element.
+  """
+  if not isinstance(dataset, data_types.DatasetV2):
+    raise TypeError(
+        f"Invalid `dataset`. Expected a `tf.data.Dataset` object "
+        f"but got {type(dataset)}.")
+
+  return dataset.get_single_element()