chore: import upstream snapshot with attribution

tensorflow/python/data/ops/BUILD

@@ -0,0 +1,297 @@
+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 = "structured_function",
+    srcs = ["structured_function.py"],
+    strict_deps = True,
+    deps = [
+        ":debug_mode",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/impl:api",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:function",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/util:function_utils",
+        "//tensorflow/python/util:lazy_loader",
+        "//tensorflow/python/util:variable_utils",
+    ],
+)
+
+py_library(
+    name = "debug_mode",
+    srcs = ["debug_mode.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "test_mode",
+    srcs = ["test_mode.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "dataset_ops",
+    # Grouped together due to mutual dependencies, to avoid dependency cycles.
+    srcs = [
+        "batch_op.py",
+        "cache_op.py",
+        "choose_from_datasets_op.py",
+        "concatenate_op.py",
+        "counter_op.py",
+        "dataset_autograph.py",
+        "dataset_ops.py",
+        "directed_interleave_op.py",
+        "filter_op.py",
+        "flat_map_op.py",
+        "from_generator_op.py",
+        "from_sparse_tensor_slices_op.py",
+        "from_tensor_slices_op.py",
+        "from_tensors_op.py",
+        "group_by_window_op.py",
+        "ignore_errors_op.py",
+        "interleave_op.py",
+        "load_op.py",
+        "map_op.py",
+        "padded_batch_op.py",
+        "prefetch_op.py",
+        "ragged_batch_op.py",
+        "random_op.py",
+        "range_op.py",
+        "rebatch_op.py",
+        "repeat_op.py",
+        "sample_from_datasets_op.py",
+        "save_op.py",
+        "scan_op.py",
+        "shard_op.py",
+        "shuffle_op.py",
+        "skip_op.py",
+        "snapshot_op.py",
+        "sparse_batch_op.py",
+        "take_op.py",
+        "take_while_op.py",
+        "unbatch_op.py",
+        "unique_op.py",
+        "window_op.py",
+        "zip_op.py",
+    ],
+    lazy_imports = True,
+    strict_deps = True,
+    deps = [
+        ":debug_mode",
+        ":iterator_ops",
+        ":options",
+        ":structured_function",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/framework:graph_proto_py_proto",
+        "//tensorflow/core/protobuf:for_core_protos_py_proto",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/autograph/operators:control_flow",
+        "//tensorflow/python/autograph/operators:py_builtins",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/experimental/ops:take_while_ops",
+        "//tensorflow/python/data/experimental/service:_pywrap_snapshot_utils",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:random_seed",
+        "//tensorflow/python/data/util:sparse",  # build_cleaner: keep
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/data/util:traverse",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:wrap_function",
+        "//tensorflow/python/framework:auto_control_deps",
+        "//tensorflow/python/framework:auto_control_deps_utils",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:function",
+        "//tensorflow/python/framework:none_tensor",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:smart_cond",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:io_ops_gen",
+        "//tensorflow/python/ops:logging_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:parsing_ops_gen",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/ops:stateless_random_ops_gen",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:string_ops_gen",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/trackable:asset",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/trackable:resource",
+        "//tensorflow/python/types:data",
+        "//tensorflow/python/types:trace",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:lazy_loader",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:numpy_compat",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+        "@absl_py//absl/logging",
+    ],
+)
+
+py_library(
+    name = "iterator_ops",
+    # Grouped together due to mutual dependencies, to avoid dependency cycles.
+    srcs = [
+        "iterator_autograph.py",
+        "iterator_ops.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":optional_ops",
+        ":options",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/operators:control_flow",
+        "//tensorflow/python/autograph/operators:py_builtins",
+        "//tensorflow/python/checkpoint:saveable_compat",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/framework:type_utils",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:parsing_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops/ragged:ragged_string_ops",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/training:saver",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "multi_device_iterator_ops",
+    srcs = ["multi_device_iterator_ops.py"],
+    strict_deps = True,
+    deps = [
+        ":dataset_ops",
+        ":iterator_ops",
+        ":options",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/framework:type_utils",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:functional_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+    ],
+)
+
+py_library(
+    name = "optional_ops",
+    srcs = ["optional_ops.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/util:structure",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:optional_ops_gen",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "options",
+    srcs = ["options.py"],
+    strict_deps = True,
+    deps = [
+        ":test_mode",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/data/util:options",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "@absl_py//absl/logging",
+    ],
+)
+
+py_library(
+    name = "readers",
+    srcs = ["readers.py"],
+    strict_deps = True,
+    deps = [
+        ":dataset_ops",
+        ":structured_function",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/util:convert",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/types:data",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)

tensorflow/python/data/ops/batch_op.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.batch`."""
+
+import warnings
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _batch(input_dataset,
+           batch_size,
+           drop_remainder=False,
+           num_parallel_calls=None,
+           deterministic=None,
+           name=None):
+  """See `Dataset.batch` for details."""
+  if num_parallel_calls is None or debug_mode.DEBUG_MODE:
+    if deterministic is not None and not debug_mode.DEBUG_MODE:
+      warnings.warn("The `deterministic` argument has no effect unless the "
+                    "`num_parallel_calls` argument is specified.")
+    return _BatchDataset(input_dataset, batch_size, drop_remainder, name=name)
+  else:
+    return _ParallelBatchDataset(
+        input_dataset,
+        batch_size,
+        drop_remainder,
+        num_parallel_calls,
+        deterministic,
+        name=name)
+
+
+class _BatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that batches contiguous elements from its input."""
+
+  def __init__(self, input_dataset, batch_size, drop_remainder, name=None):
+    """See `Dataset.batch()` for details."""
+    self._input_dataset = input_dataset
+    self._batch_size = ops.convert_to_tensor(
+        batch_size, dtype=dtypes.int64, name="batch_size")
+    self._drop_remainder = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+
+    constant_drop_remainder = tensor_util.constant_value(self._drop_remainder)
+    # 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
+      constant_batch_size = tensor_util.constant_value(self._batch_size)
+      self._structure = nest.map_structure(
+          lambda component_spec: component_spec._batch(constant_batch_size),
+          input_dataset.element_spec)
+    else:
+      self._structure = nest.map_structure(
+          lambda component_spec: component_spec._batch(None),
+          input_dataset.element_spec)
+
+    self._name = name
+    variant_tensor = gen_dataset_ops.batch_dataset_v2(
+        input_dataset._variant_tensor,
+        batch_size=self._batch_size,
+        drop_remainder=self._drop_remainder,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+
+class _ParallelBatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that batches contiguous elements from its input in parallel."""
+
+  def __init__(self,
+               input_dataset,
+               batch_size,
+               drop_remainder,
+               num_parallel_calls,
+               deterministic,
+               name=None):
+    """See `Dataset.batch()` for details."""
+    self._input_dataset = input_dataset
+    self._batch_size = ops.convert_to_tensor(
+        batch_size, dtype=dtypes.int64, name="batch_size")
+    self._drop_remainder = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+    self._num_parallel_calls = ops.convert_to_tensor(
+        num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
+    if deterministic is None:
+      self._deterministic = "default"
+    elif deterministic:
+      self._deterministic = "true"
+    else:
+      self._deterministic = "false"
+
+    constant_drop_remainder = tensor_util.constant_value(self._drop_remainder)
+    # 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
+      constant_batch_size = tensor_util.constant_value(self._batch_size)
+      self._structure = nest.map_structure(
+          lambda component_spec: component_spec._batch(constant_batch_size),
+          input_dataset.element_spec)
+    else:
+      self._structure = nest.map_structure(
+          lambda component_spec: component_spec._batch(None),
+          input_dataset.element_spec)
+
+    self._name = name
+    variant_tensor = gen_dataset_ops.parallel_batch_dataset(
+        input_dataset._variant_tensor,
+        batch_size=self._batch_size,
+        num_parallel_calls=self._num_parallel_calls,
+        drop_remainder=self._drop_remainder,
+        deterministic=self._deterministic,
+        **self._common_args)
+
+    super().__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/cache_op.py

@@ -0,0 +1,49 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.cache`."""
+
+from tensorflow.python import tf2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _cache(input_dataset, filename, name):  # pylint: disable=unused-private-name
+  return CacheDataset(input_dataset, filename, name)
+
+
+class CacheDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that caches elements of its input."""
+
+  def __init__(self, input_dataset, filename, name=None):
+    """See `Dataset.cache()` for details."""
+    self._input_dataset = input_dataset
+    self._filename = ops.convert_to_tensor(
+        filename, dtype=dtypes.string, name="filename")
+    self._name = name
+    if tf2.enabled() and (context.executing_eagerly() or ops.inside_function()):
+      variant_tensor = gen_dataset_ops.cache_dataset_v2(
+          input_dataset._variant_tensor,  # pylint: disable=protected-access
+          filename=self._filename,
+          cache=gen_dataset_ops.dummy_memory_cache(),
+          **self._common_args)
+    else:
+      variant_tensor = gen_dataset_ops.cache_dataset(
+          input_dataset._variant_tensor,  # pylint: disable=protected-access
+          filename=self._filename,
+          **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/choose_from_datasets_op.py

@@ -0,0 +1,53 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.choose_from_datasets`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import directed_interleave_op
+from tensorflow.python.data.util import structure
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.types import data as data_types
+
+
+def _choose_from_datasets(  # pylint: disable=unused-private-name
+    datasets, choice_dataset, stop_on_empty_dataset=True
+):
+  """See `Dataset.choose_from_datasets()` for details."""
+
+  if not datasets:
+    raise ValueError("Invalid `datasets`. `datasets` should not be empty.")
+  if not isinstance(choice_dataset, data_types.DatasetV2):
+    raise TypeError(
+        "Invalid `choice_dataset`. `choice_dataset` should be a "
+        f"`tf.data.Dataset` but is {type(choice_dataset)}."
+    )
+  if not structure.are_compatible(
+      choice_dataset.element_spec, tensor_spec.TensorSpec([], dtypes.int64)
+  ):
+    raise TypeError(
+        "Invalid `choice_dataset`. Elements of `choice_dataset` "
+        "must be scalar `tf.int64` tensors but are "
+        f"{choice_dataset.element_spec}."
+    )
+  # Replicates the `choice_dataset` component so that each split makes choices
+  # independently. This avoids the need for prohibitively expensive
+  # cross-split coordination.
+  choice_dataset = dataset_ops.apply_rewrite(
+      choice_dataset, "replicate_on_split"
+  )
+  return directed_interleave_op._directed_interleave(  # pylint: disable=protected-access
+      choice_dataset, datasets, stop_on_empty_dataset
+  )

tensorflow/python/data/ops/concatenate_op.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.concatenate`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.util import nest as tf_nest
+
+
+def _concatenate(input_dataset, dataset_to_concatenate, name):  # pylint: disable=unused-private-name
+  return _ConcatenateDataset(input_dataset, dataset_to_concatenate, name)
+
+
+class _ConcatenateDataset(dataset_ops.DatasetV2):
+  """A `Dataset` that concatenates its input with given dataset."""
+
+  def __init__(self, input_dataset, dataset_to_concatenate, name=None):
+    """See `Dataset.concatenate()` for details."""
+    self._input_dataset = input_dataset
+    self._dataset_to_concatenate = dataset_to_concatenate
+
+    def common_supertype(a, b):
+      result = a.most_specific_common_supertype([b])
+      if result is None:
+        raise TypeError(f"No common supertype of {a} and {b}.")
+      return result
+
+    try:
+      self._structure = tf_nest.map_structure(
+          common_supertype, input_dataset.element_spec,
+          dataset_to_concatenate.element_spec)
+    except (TypeError, ValueError) as e:
+      raise TypeError(f"Incompatible dataset elements:\n"
+                      f"  {input_dataset.element_spec} vs. "
+                      f"  {dataset_to_concatenate.element_spec}") from e
+
+    self._input_datasets = [input_dataset, dataset_to_concatenate]
+    self._name = name
+    # pylint: disable=protected-access
+    variant_tensor = gen_dataset_ops.concatenate_dataset(
+        input_dataset._variant_tensor, dataset_to_concatenate._variant_tensor,
+        **self._common_args)
+    # pylint: enable=protected-access
+    super().__init__(variant_tensor)
+
+  def _inputs(self):
+    return self._input_datasets
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/counter_op.py

@@ -0,0 +1,26 @@
+# 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 implementation of `tf.data.Dataset.counter`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import ops
+
+
+def _counter(start, step, dtype, name=None):
+  with ops.name_scope("counter"):
+    start = ops.convert_to_tensor(start, dtype=dtype, name="start")
+    step = ops.convert_to_tensor(step, dtype=dtype, name="step")
+    return (dataset_ops.Dataset.from_tensors(0, name=name).repeat(None).scan(
+        start, lambda state, _: (state + step, state)))

tensorflow/python/data/ops/dataset_autograph.py

@@ -0,0 +1,224 @@
+# 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.
+# ==============================================================================
+"""Autograph specific overrides for dataset_ops."""
+from tensorflow.python.autograph.operators import control_flow
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.data.experimental.ops import take_while_ops
+from tensorflow.python.data.ops import dataset_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 array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import gen_string_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util import nest
+
+
+def _general_purpose_scan(ds, init_state, body):
+  """Variant of Dataset.scan with semantics of general-purpose computation."""
+  # Datasets are typically intended for data preprocessing. However, in
+  # autograph loops they usually appear as general-purpose computations (for
+  # example, a custom training loop). These two use cases require significantly
+  # different optimization policies, the most important of which is the device
+  # placement. The flag override for use_default_device below instructs the
+  # runtime to treat the computation as general-purpose, rather than data
+  # preprocessing.
+
+  # Loaded lazily due to a circular dependency (dataset_ops ->
+  # scan_op -> dataset_ops).
+  # pylint: disable=g-import-not-at-top,protected-access
+  from tensorflow.python.data.ops import scan_op
+  return scan_op._ScanDataset(ds, init_state, body, use_default_device=False)
+  # pylint: enable=g-import-not-at-top,protected-access
+
+
+def _tf_ag_dataset_for_stmt(
+    ds, extra_test, body, get_state, set_state, symbol_names, opts
+):
+  """Overload of _dataset_for_stmt with early stopping. See for_stmt."""
+  # Note: This is easier to follow with the insight that the computations in
+  # a dataset pipeline are transposed (aka fused).
+  # For example, given a pipeline input -> scan -> take_while -> reduce,
+  # and a dataset with input [1, 2, 3], the computations occur in the following
+  # order:
+  #  reduce(take_while(scan(1)))
+  #  reduce(take_while(scan(2)))
+  #  reduce(take_while(scan(3)))
+
+  init_vars = get_state()
+  control_flow.verify_loop_init_vars(init_vars, symbol_names)
+
+  # Workaround for Dataset.reduce not allowing empty state tensors - create
+  # a dummy state variable that remains unused.
+  # TODO(mdan): reduce should allow and match empty structures.
+  if not init_vars:
+    init_vars = (constant_op.constant(0),)
+    symbol_names = ("<internal dummy>",)
+
+    def dummy_set_state(unused_dummy):
+      pass
+
+    def dummy_get_state():
+      return (constant_op.constant(0),)
+
+    get_state, set_state = dummy_get_state, dummy_set_state
+
+  def scan_body(scan_state, scan_inputs):
+    """Main body of the Dataset.scan."""
+    loop_vars, iterate = scan_state, scan_inputs
+    set_state(loop_vars)
+
+    def main_path():
+      body(iterate)
+      new_loop_vars = get_state()
+      control_flow.verify_tf_loop_vars(
+          init_vars,
+          loop_vars,
+          new_loop_vars,
+          symbol_names,
+          opts,
+          check_shapes=False)
+      return new_loop_vars
+
+    if extra_test is not None:
+      extra_cond = extra_test()
+      new_loop_vars = cond.cond(extra_cond, main_path,
+                                lambda: loop_vars)
+    else:
+      # TODO(mdan): the optimizer should be able to remove an invariant cond?
+      extra_cond = (constant_op.constant(True),)  # dummy value, unused
+      new_loop_vars = main_path()
+
+    scan_outputs = new_loop_vars, extra_cond
+    new_scan_state = new_loop_vars
+    return new_scan_state, scan_outputs
+
+  def take_while_predicate(unused_loop_vars, extra_cond):
+    return extra_cond
+
+  def reduce_body(unused_reduce_state, scan_outputs):
+    output_loop_vars, unused_extra_cond = scan_outputs
+    new_reduce_state = output_loop_vars
+    return new_reduce_state
+
+  ds = _general_purpose_scan(ds, init_vars, scan_body)
+  if extra_test is not None:
+    ds = ds.apply(take_while_ops.take_while(take_while_predicate))
+  final_loop_vars = ds.reduce(init_vars, reduce_body)
+  set_state(final_loop_vars)
+
+
+def _tf_ag_dataset_abs(ds):
+  specs = nest.flatten(ds.element_spec)
+  if len(specs) == 1:
+    return ds.map(math_ops.abs, num_parallel_calls=dataset_ops.AUTOTUNE)
+  return ds.map(
+      lambda *e: nest.map_structure(math_ops.abs, e),
+      num_parallel_calls=dataset_ops.AUTOTUNE)
+
+
+def _tf_ag_dataset_len(s):
+  """Autograph override of the builtin len for dataset_ops.DataSetV2."""
+  l = s.cardinality()
+  msg = gen_string_ops.string_join([
+      "len requires dataset with definitive cardinality, got ",
+      gen_string_ops.as_string(l),
+  ])
+  # TODO(yongtang): UNKNOWN is treated as an error.
+  # In case there are more UNKNOWN cases for dataset, we could
+  # use dataset.reduce() to find out the length (in an expensive way).
+  with ops.control_dependencies([
+      control_flow_assert.Assert(
+          math_ops.logical_and(
+              math_ops.not_equal(l, dataset_ops.INFINITE),
+              math_ops.not_equal(l, dataset_ops.UNKNOWN)), [msg])
+  ]):
+    l = array_ops.identity(l)
+
+  return l
+
+
+def _tf_ag_dataset_enumerate(ds, start=0):
+  return ds.enumerate(start)
+
+
+def _tf_ag_dataset_zip(*iterables, strict=False):
+  if strict:
+    raise ValueError("strict zip not supported by Dataset")
+  return dataset_ops.DatasetV2.zip(iterables)
+
+
+def _tf_ag_dataset_map(fn, *iterables):
+  return dataset_ops.DatasetV2.zip(iterables).map(fn)
+
+
+def _tf_ag_dataset_filter(fn, iterable):
+  return iterable.filter(fn)
+
+
+# any() operation is essentially a "if first True element exist".
+# For that it could be translated to `filter(True)` to filter out
+# only `True` element, and then `take(1)`. This works in tf.data
+# as tf.data's filter+take is done in pipeline so it will stop
+# as soon as `take(1)` returns.
+def _tf_ag_dataset_any(iterable):
+  # check and make sure iterable.element_spec only consists of one
+  # element of tf.bool.
+  specs = nest.flatten(iterable.element_spec)
+  if len(specs) != 1 or specs[0].dtype != dtypes.bool:
+    raise ValueError('in graph mode, the "any" builtin only supports datasets '
+                     'that return bool scalars; got: {}'.format(
+                         iterable.element_spec))
+  ds = iterable.filter(lambda x: x)
+  ds = ds.take(1)
+  ds = ds.reduce(constant_op.constant(False, dtype=dtypes.bool), lambda _, y: y)
+  return ds
+
+
+# all() operation is similar to any() and could be translated
+# to `filter(False)` then `take(1)`, and check if `False` exists.
+def _tf_ag_dataset_all(iterable):
+  # check and make sure iterable.element_spec only consists of one
+  # element of tf.bool.
+  specs = nest.flatten(iterable.element_spec)
+  if len(specs) != 1 or specs[0].dtype != dtypes.bool:
+    raise ValueError('in graph mode, the "all" builtin only supports datasets '
+                     'that return bool scalars; got: {}'.format(
+                         iterable.element_spec))
+  ds = iterable.filter(math_ops.logical_not)
+  ds = ds.take(1)
+  ds = ds.reduce(constant_op.constant(True, dtype=dtypes.bool), lambda _, y: y)
+  return ds
+
+
+def register_overrides():
+  """Registers the autograph specific overrides for dataset_ops."""
+  control_flow.for_loop_registry.register(
+      dataset_ops.DatasetV2, _tf_ag_dataset_for_stmt
+  )
+  py_builtins.abs_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_abs)
+  py_builtins.len_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_len)
+  py_builtins.enumerate_registry.register(
+      dataset_ops.DatasetV2, _tf_ag_dataset_enumerate
+  )
+  py_builtins.zip_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_zip)
+  py_builtins.map_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_map)
+  py_builtins.filter_registry.register(
+      dataset_ops.DatasetV2, _tf_ag_dataset_filter
+  )
+  py_builtins.any_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_any)
+  py_builtins.all_registry.register(dataset_ops.DatasetV2, _tf_ag_dataset_all)

tensorflow/python/data/ops/debug_mode.py

@@ -0,0 +1,77 @@
+# 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 debug mode enabler."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.util.tf_export import tf_export
+
+
+DEBUG_MODE = False
+
+
+@tf_export("data.experimental.enable_debug_mode")
+def enable_debug_mode():
+  """Enables debug mode for tf.data.
+
+  Example usage with pdb module:
+  ```
+  import tensorflow as tf
+  import pdb
+
+  tf.data.experimental.enable_debug_mode()
+
+  def func(x):
+    # Python 3.7 and older requires `pdb.Pdb(nosigint=True).set_trace()`
+    pdb.set_trace()
+    x = x + 1
+    return x
+
+  dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
+  dataset = dataset.map(func)
+
+  for item in dataset:
+    print(item)
+  ```
+
+  The effect of debug mode is two-fold:
+
+  1) Any transformations that would introduce asynchrony, parallelism, or
+  non-determinism to the input pipeline execution will be forced to execute
+  synchronously, sequentially, and deterministically.
+
+  2) Any user-defined functions passed into tf.data transformations such as
+  `map` will be wrapped in `tf.py_function` so that their body is executed
+  "eagerly" as a Python function as opposed to a traced TensorFlow graph, which
+  is the default behavior. Note that even when debug mode is enabled, the
+  user-defined function is still traced  to infer the shape and type of its
+  outputs; as a consequence, any `print` statements or breakpoints will be
+  triggered once during the tracing before the actual execution of the input
+  pipeline.
+
+  NOTE: As the debug mode setting affects the construction of the tf.data input
+  pipeline, it should be enabled before any tf.data definitions.
+
+  Raises:
+    ValueError: When invoked from graph mode.
+  """
+  if context.executing_eagerly():
+    toggle_debug_mode(True)
+  else:
+    raise ValueError("`enable_debug_mode() is only supported in eager mode.")
+
+
+def toggle_debug_mode(debug_mode):
+  global DEBUG_MODE
+  DEBUG_MODE = debug_mode

tensorflow/python/data/ops/directed_interleave_op.py

@@ -0,0 +1,72 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.shuffle`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _directed_interleave(  # pylint: disable=unused-private-name
+    selector_input, data_inputs, stop_on_empty_dataset=False
+):
+  return _DirectedInterleaveDataset(
+      selector_input, data_inputs, stop_on_empty_dataset=stop_on_empty_dataset
+  )
+
+
+class _DirectedInterleaveDataset(dataset_ops.DatasetV2):
+  """A substitute for `Dataset.interleave()` on a fixed list of datasets."""
+
+  def __init__(self, selector_input, data_inputs, stop_on_empty_dataset=False):
+    self._selector_input = selector_input
+    self._data_inputs = list(data_inputs)
+    self._stop_on_empty_dataset = stop_on_empty_dataset
+
+    spec = self._data_inputs[0].element_spec
+    for i, data_input in enumerate(self._data_inputs[1:]):
+      def common_supertype(a, b):
+        result = a.most_specific_common_supertype([b])
+        if result is None:
+          raise TypeError(f"No common supertype of {a} and {b}.")
+        return result
+
+      try:
+        spec = nest.map_structure(common_supertype, spec,
+                                  data_input.element_spec)
+      except (TypeError, ValueError) as e:
+        raise TypeError(f"Invalid `datasets`. `datasets` must have compatible "
+                        f"element specs.\n Dataset 0 "
+                        f"element_spec={data_inputs[0].element_spec}.\n"
+                        f"Dataset {i+1} "
+                        f"element_spec={data_input.element_spec}.") from e
+    self._element_spec = spec
+
+    # pylint: disable=protected-access
+    variant_tensor = (
+        ged_ops.directed_interleave_dataset(
+            self._selector_input._variant_tensor,
+            [data_input._variant_tensor for data_input in self._data_inputs],
+            stop_on_empty_dataset=self._stop_on_empty_dataset,
+            **self._flat_structure))
+
+    super().__init__(variant_tensor)
+
+  def _inputs(self):
+    return [self._selector_input] + self._data_inputs
+
+  @property
+  def element_spec(self):
+    return self._element_spec

tensorflow/python/data/ops/filter_op.py

@@ -0,0 +1,61 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.filter`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _filter(input_dataset, predicate, name=None):  # pylint: disable=redefined-builtin
+  return _FilterDataset(input_dataset, predicate, name=name)
+
+
+class _FilterDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that filters its input according to a predicate function."""
+
+  def __init__(self,
+               input_dataset,
+               predicate,
+               use_legacy_function=False,
+               name=None):
+    """See `Dataset.filter` for details."""
+    self._input_dataset = input_dataset
+    wrapped_func = structured_function.StructuredFunctionWrapper(
+        predicate,
+        self._transformation_name(),
+        dataset=input_dataset,
+        use_legacy_function=use_legacy_function)
+    if not wrapped_func.output_structure.is_compatible_with(
+        tensor_spec.TensorSpec([], dtypes.bool)):
+      raise ValueError(f"Invalid `predicate`. `predicate` must return a "
+                       f"`tf.bool` scalar tensor, but its return type is "
+                       f"{wrapped_func.output_structure}.")
+    self._predicate = wrapped_func
+    self._name = name
+    variant_tensor = gen_dataset_ops.filter_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        other_arguments=self._predicate.function.captured_inputs,
+        predicate=self._predicate.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._predicate]
+
+  def _transformation_name(self):
+    return "Dataset.filter()"

tensorflow/python/data/ops/flat_map_op.py

@@ -0,0 +1,57 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.flat_map`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _flat_map(input_dataset, map_func, name=None):  # pylint: disable=unused-private-name
+  """See `Dataset.flat_map()` for details."""
+  return _FlatMapDataset(input_dataset, map_func, name)
+
+
+class _FlatMapDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over its input and flattens the result."""
+
+  def __init__(self, input_dataset, map_func, name=None):
+
+    self._input_dataset = input_dataset
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func, self._transformation_name(), dataset=input_dataset)
+    if not isinstance(self._map_func.output_structure, dataset_ops.DatasetSpec):
+      raise TypeError(
+          "The `map_func` argument must return a `Dataset` object. Got "
+          f"{dataset_ops.get_type(self._map_func.output_structure)!r}.")
+    # pylint: disable=protected-access
+    self._structure = self._map_func.output_structure._element_spec
+    self._name = name
+    variant_tensor = gen_dataset_ops.flat_map_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        f=self._map_func.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+  def _transformation_name(self):
+    return "Dataset.flat_map()"

tensorflow/python/data/ops/from_generator_op.py

@@ -0,0 +1,400 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.from_generator`."""
+
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+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_spec
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import script_ops
+
+
+def _from_generator(generator, output_types, output_shapes, args,
+                    output_signature, name):
+  """Creates a `Dataset` whose elements are generated by `generator`.
+
+  Note: The current implementation of `Dataset.from_generator()` uses
+  `tf.numpy_function` and inherits the same constraints. In particular, it
+  requires the dataset and iterator related operations to be placed
+  on a device in the same process as the Python program that called
+  `Dataset.from_generator()`. In particular, using `from_generator` will
+  preclude the use of tf.data service for scaling out dataset processing.
+  The body of `generator` will not be serialized in a `GraphDef`, and you
+  should not use this method if you need to serialize your model and restore
+  it in a different environment.
+
+  The `generator` argument must be a callable object that returns
+  an object that supports the `iter()` protocol (e.g. a generator function).
+
+  The elements generated by `generator` must be compatible with either the
+  given `output_signature` argument or with the given `output_types` and
+  (optionally) `output_shapes` arguments, whichever was specified.
+
+  The recommended way to call `from_generator` is to use the
+  `output_signature` argument. In this case the output will be assumed to
+  consist of objects with the classes, shapes and types defined by
+  `tf.TypeSpec` objects from `output_signature` argument:
+
+  >>> def gen():
+  ...   ragged_tensor = tf.ragged.constant([[1, 2], [3]])
+  ...   yield 42, ragged_tensor
+  >>>
+  >>> dataset = tf.data.Dataset.from_generator(
+  ...      gen,
+  ...      output_signature=(
+  ...          tf.TensorSpec(shape=(), dtype=tf.int32),
+  ...          tf.RaggedTensorSpec(shape=(2, None), dtype=tf.int32)))
+  >>>
+  >>> list(dataset.take(1))
+  [(<tf.Tensor: shape=(), dtype=int32, numpy=42>,
+  <tf.RaggedTensor [[1, 2], [3]]>)]
+
+  There is also a deprecated way to call `from_generator` by either with
+  `output_types` argument alone or together with `output_shapes` argument.
+  In this case the output of the function will be assumed to consist of
+  `tf.Tensor` objects with the types defined by `output_types` and with the
+  shapes which are either unknown or defined by `output_shapes`.
+
+  Note: If `generator` depends on mutable global variables or other external
+  state, be aware that the runtime may invoke `generator` multiple times
+  (in order to support repeating the `Dataset`) and at any time
+  between the call to `Dataset.from_generator()` and the production of the
+  first element from the generator. Mutating global variables or external
+  state can cause undefined behavior, and we recommend that you explicitly
+  cache any external state in `generator` before calling
+  `Dataset.from_generator()`.
+
+  Note: While the `output_signature` parameter makes it possible to yield
+  `Dataset` elements, the scope of `Dataset.from_generator()` should be
+  limited to logic that cannot be expressed through tf.data operations. Using
+  tf.data operations within the generator function is an anti-pattern and may
+  result in incremental memory growth.
+
+  Args:
+    generator: A callable object that returns an object that supports the
+      `iter()` protocol. If `args` is not specified, `generator` must take no
+      arguments; otherwise it must take as many arguments as there are values in
+      `args`.
+    output_types: (Optional.) A (nested) structure of `tf.DType` objects
+      corresponding to each component of an element yielded by `generator`.
+    output_shapes: (Optional.) A (nested) structure of `tf.TensorShape` objects
+      corresponding to each component of an element yielded by `generator`.
+    args: (Optional.) A tuple of `tf.Tensor` objects that will be evaluated and
+      passed to `generator` as NumPy-array arguments.
+    output_signature: (Optional.) A (nested) structure of `tf.TypeSpec` objects
+      corresponding to each component of an element yielded by `generator`.
+    name: (Optional.) A name for the tf.data operations used by
+      `from_generator`.
+
+  Returns:
+    Dataset: A `Dataset`.
+  """
+  if not callable(generator):
+    raise TypeError("`generator` must be a Python callable.")
+
+  if output_signature is not None:
+    if output_types is not None:
+      raise TypeError("The `output_types` argument can not be used together "
+                      "with the `output_signature` argument.")
+    if output_shapes is not None:
+      raise TypeError("The `output_shapes` argument can not be used together "
+                      "with the `output_signature` argument.")
+    for spec in nest.flatten(output_signature):
+      if not isinstance(spec, type_spec.TypeSpec):
+        raise TypeError(f"`output_signature` must contain objects that are "
+                        f"subclass of `tf.TypeSpec` but found {type(spec)} "
+                        f"which is not.")
+  else:
+    if output_types is None:
+      raise TypeError("To specify the output signature you need to provide "
+                      "either the `output_signature` argument or the "
+                      "`output_types` argument.")
+
+  if output_signature is None:
+    if output_shapes is None:
+      output_shapes = nest.map_structure(
+          lambda _: tensor_shape.TensorShape(None), output_types)
+    else:
+      output_shapes = nest.map_structure_up_to(output_types,
+                                               tensor_shape.as_shape,
+                                               output_shapes)
+    output_signature = nest.map_structure_up_to(output_types,
+                                                tensor_spec.TensorSpec,
+                                                output_shapes, output_types)
+  if all(
+      isinstance(x, tensor_spec.TensorSpec)
+      for x in nest.flatten(output_signature)):
+    output_types = nest.pack_sequence_as(
+        output_signature, [x.dtype for x in nest.flatten(output_signature)])
+    output_shapes = nest.pack_sequence_as(
+        output_signature, [x.shape for x in nest.flatten(output_signature)])
+
+  if args is None:
+    args = ()
+  else:
+    args = tuple(ops.convert_n_to_tensor(args, name="args"))
+
+  generator_state = dataset_ops.DatasetV2._GeneratorState(generator)  # pylint: disable=protected-access
+
+  def get_iterator_id_fn(unused_dummy):
+    """Creates a unique `iterator_id` for each pass over the dataset.
+
+    The returned `iterator_id` disambiguates between multiple concurrently
+    existing iterators.
+
+    Args:
+      unused_dummy: Ignored value.
+
+    Returns:
+      A `tf.int64` tensor whose value uniquely identifies an iterator in
+      `generator_state`.
+    """
+    return script_ops.numpy_function(generator_state.get_next_id, args,
+                                     dtypes.int64)
+
+  def generator_next_fn(iterator_id_t):
+    """Generates the next element from iterator with ID `iterator_id_t`.
+
+    We map this function across an infinite repetition of the
+    `iterator_id_t`, and raise `StopIteration` to terminate the iteration.
+
+    Args:
+      iterator_id_t: A `tf.int64` tensor whose value uniquely identifies the
+        iterator in `generator_state` from which to generate an element.
+
+    Returns:
+      The next element to generate from the iterator.
+    """
+    if output_types and output_shapes:
+      flattened_types = [
+          dtypes.as_dtype(dt) for dt in nest.flatten(output_types)
+      ]
+      flattened_shapes = nest.flatten(output_shapes)
+
+      def generator_py_func(iterator_id):
+        """A `py_func` that will be called to invoke the iterator."""
+        # `next()` raises `StopIteration` when there are no more
+        # elements remaining to be generated.
+        values = next(generator_state.get_iterator(iterator_id))
+
+        # Use the same _convert function from the py_func() implementation to
+        # convert the returned values to arrays early, so that we can inspect
+        # their values.
+        try:
+          flattened_values = nest.flatten_up_to(output_types, values)
+        except (TypeError, ValueError) as e:
+          raise TypeError(
+              f"`generator` yielded an element that did not match the "
+              f"expected structure. The expected structure was "
+              f"{output_types}, but the yielded element was {values}.") from e
+        ret_arrays = []
+        for ret, dtype in zip(flattened_values, flattened_types):
+          try:
+            ret_arrays.append(
+                script_ops.FuncRegistry._convert(  # pylint: disable=protected-access
+                    ret,
+                    dtype=dtype.as_numpy_dtype))
+          except (TypeError, ValueError) as e:
+            raise TypeError(
+                f"`generator` yielded an element that could not be "
+                f"converted to the expected type. The expected type was "
+                f"{dtype.name}, but the yielded element was {ret}.") from e
+
+        # Additional type and shape checking to ensure that the components of
+        # the generated element match the `output_types` and `output_shapes`
+        # arguments.
+        for (ret_array, expected_dtype,
+             expected_shape) in zip(ret_arrays, flattened_types,
+                                    flattened_shapes):
+          if ret_array.dtype != expected_dtype.as_numpy_dtype:
+            raise TypeError(
+                f"`generator` yielded an element of type {ret_array.dtype} "
+                f"where an element of type {expected_dtype.as_numpy_dtype} "
+                f"was expected.")
+          if not expected_shape.is_compatible_with(ret_array.shape):
+            raise TypeError(
+                f"`generator` yielded an element of shape {ret_array.shape} "
+                f"where an element of shape {expected_shape} was expected.")
+
+        return ret_arrays
+
+      flat_values = script_ops.numpy_function(generator_py_func,
+                                              [iterator_id_t], flattened_types)
+
+      # In debug mode the numpy_function will return a scalar if
+      # generator_py_func produces only a single value.
+      if not isinstance(flat_values, (list, tuple)):
+        flat_values = [flat_values]
+
+      # The `py_func()` op drops the inferred shapes, so we add them back in
+      # here.
+      if output_shapes is not None:
+        for ret_t, shape in zip(flat_values, flattened_shapes):
+          ret_t.set_shape(shape)
+
+      return nest.pack_sequence_as(output_types, flat_values)
+    else:
+      flat_output_types = structure.get_flat_tensor_types(output_signature)
+
+      def generator_py_func(iterator_id):
+        """A `py_func` that will be called to invoke the iterator."""
+        # `next()` raises `StopIteration` when there are no more
+        # elements remaining to be generated.
+        values = next(generator_state.get_iterator(iterator_id.numpy()))
+
+        try:
+          values = structure.normalize_element(values, output_signature)
+        except (TypeError, ValueError) as e:
+          raise TypeError(
+              f"`generator` yielded an element that did not match the "
+              f"expected structure. The expected structure was "
+              f"{output_signature}, but the yielded element was "
+              f"{values}.") from e
+
+        values_spec = structure.type_spec_from_value(values)
+
+        if not structure.are_compatible(values_spec, output_signature):
+          raise TypeError(
+              f"`generator` yielded an element of {values_spec} where an "
+              f"element of {output_signature} was expected.")
+
+        return structure.to_tensor_list(output_signature, values)
+
+      return script_ops.eager_py_func(
+          generator_py_func, inp=[iterator_id_t], Tout=flat_output_types)
+
+  def finalize_fn(iterator_id_t):
+    """Releases host-side state for the iterator with ID `iterator_id_t`."""
+
+    def finalize_py_func(iterator_id):
+      generator_state.iterator_completed(iterator_id)
+      # We return a dummy value so that the `finalize_fn` has a valid
+      # signature.
+      # NOTE(mrry): Explicitly create an array of `np.int64` because implicit
+      # casting in `py_func()` will create an array of `np.int32` on Windows,
+      # leading to a runtime error.
+      return np.array(0, dtype=np.int64)
+
+    return script_ops.numpy_function(finalize_py_func, [iterator_id_t],
+                                     dtypes.int64)
+
+  # This function associates each traversal of `generator` with a unique
+  # iterator ID.
+  def flat_map_fn(dummy_arg):
+    # The `get_iterator_id_fn` gets a unique ID for the current instance of
+    # of the generator.
+    # The `generator_next_fn` gets the next element from the iterator with the
+    # given ID, and raises StopIteration when that iterator contains no
+    # more elements.
+    return _GeneratorDataset(
+        dummy_arg,
+        get_iterator_id_fn,
+        generator_next_fn,
+        finalize_fn,
+        output_signature,
+        name=name)
+
+  # A single-element dataset that, each time it is evaluated, contains a
+  # freshly-generated and unique (for the returned dataset) int64
+  # ID that will be used to identify the appropriate Python state, which
+  # is encapsulated in `generator_state`, and captured in
+  # `get_iterator_id_map_fn`.
+  dummy = 0
+  id_dataset = dataset_ops.Dataset.from_tensors(dummy, name=name)
+
+  # A dataset that contains all of the elements generated by a
+  # single iterator created from `generator`, identified by the
+  # iterator ID contained in `id_dataset`. Lifting the iteration
+  # into a flat_map here enables multiple repetitions and/or nested
+  # versions of the returned dataset to be created, because it forces
+  # the generation of a new ID for each version.
+  return id_dataset.flat_map(flat_map_fn, name=name)
+
+
+class _GeneratorDataset(dataset_ops.DatasetSource):
+  """A `Dataset` that generates elements by invoking a function."""
+
+  def __init__(self,
+               init_args,
+               init_func,
+               next_func,
+               finalize_func,
+               output_signature,
+               name=None):
+    """Constructs a `_GeneratorDataset`.
+
+    Args:
+      init_args: A (nested) structure representing the arguments to `init_func`.
+      init_func: A TensorFlow function that will be called on `init_args` each
+        time a C++ iterator over this dataset is constructed. Returns a (nested)
+        structure representing the "state" of the dataset.
+      next_func: A TensorFlow function that will be called on the result of
+        `init_func` to produce each element, and that raises `OutOfRangeError`
+        to terminate iteration.
+      finalize_func: A TensorFlow function that will be called on the result of
+        `init_func` immediately before a C++ iterator over this dataset is
+        destroyed. The return value is ignored.
+      output_signature: A (nested) structure of `tf.TypeSpec` objects describing
+        the output of `next_func`.
+      name: Optional. A name for the tf.data transformation.
+    """
+    self._init_args = init_args
+
+    self._init_structure = structure.type_spec_from_value(init_args)
+
+    self._init_func = structured_function.StructuredFunctionWrapper(
+        init_func,
+        self._transformation_name(),
+        input_structure=self._init_structure)
+
+    self._next_func = structured_function.StructuredFunctionWrapper(
+        next_func,
+        self._transformation_name(),
+        input_structure=self._init_func.output_structure)
+
+    self._finalize_func = structured_function.StructuredFunctionWrapper(
+        finalize_func,
+        self._transformation_name(),
+        input_structure=self._init_func.output_structure)
+
+    self._output_signature = output_signature
+
+    self._name = name
+
+    variant_tensor = gen_dataset_ops.generator_dataset(
+        structure.to_tensor_list(self._init_structure, self._init_args) +
+        self._init_func.function.captured_inputs,
+        self._next_func.function.captured_inputs,
+        self._finalize_func.function.captured_inputs,
+        init_func=self._init_func.function,
+        next_func=self._next_func.function,
+        finalize_func=self._finalize_func.function,
+        **self._common_args)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._output_signature
+
+  def _transformation_name(self):
+    return "Dataset.from_generator()"

tensorflow/python/data/ops/from_sparse_tensor_slices_op.py

@@ -0,0 +1,53 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.from_sparse_tensor_slices`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _from_sparse_tensor_slices(sparse_tensor):  # pylint: disable=unused-private-name
+  return dataset_ops.DatasetV1Adapter(_SparseTensorSliceDataset(sparse_tensor))
+
+
+class _SparseTensorSliceDataset(dataset_ops.DatasetSource):
+  """A `Dataset` that splits a rank-N `tf.sparse.SparseTensor` into its rows."""
+
+  def __init__(self, sparse_tensor):
+    """See `Dataset.from_sparse_tensor_slices()` for details."""
+    if not isinstance(sparse_tensor, sparse_tensor_lib.SparseTensor):
+      raise TypeError(f"Invalid `sparse_tensor`. `sparse_tensor` must be a "
+                      f"`tf.sparse.SparseTensor`. Got {type(sparse_tensor)}.")
+    self._sparse_tensor = sparse_tensor
+
+    indices_shape = self._sparse_tensor.indices.get_shape()
+    shape_shape = self._sparse_tensor.dense_shape.get_shape()
+    rank = (indices_shape.dims[1] - 1).merge_with(shape_shape.dims[0] - 1)
+    self._structure = (tensor_spec.TensorSpec([None, rank], dtypes.int64),
+                       tensor_spec.TensorSpec([None],
+                                              self._sparse_tensor.dtype),
+                       tensor_spec.TensorSpec([rank], dtypes.int64))
+
+    variant_tensor = gen_dataset_ops.sparse_tensor_slice_dataset(
+        self._sparse_tensor.indices, self._sparse_tensor.values,
+        self._sparse_tensor.dense_shape)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/from_tensor_slices_op.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.from_tensor_slices`."""
+
+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.framework import tensor_shape
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _from_tensor_slices(tensors, name=None):
+  return _TensorSliceDataset(tensors, name=name)
+
+
+class _TensorSliceDataset(dataset_ops.DatasetSource):
+  """A `Dataset` of slices from a dataset element."""
+
+  def __init__(self, element, is_files=False, name=None):
+    """See `Dataset.from_tensor_slices` for details."""
+    element = structure.normalize_element(element)
+    batched_spec = structure.type_spec_from_value(element)
+    self._tensors = structure.to_batched_tensor_list(batched_spec, element)
+    if not self._tensors:
+      raise ValueError("Invalid `element`. `element` should not be empty.")
+    self._structure = nest.map_structure(
+        lambda component_spec: component_spec._unbatch(), batched_spec)  # pylint: disable=protected-access
+    self._name = name
+
+    batch_dim = tensor_shape.Dimension(
+        tensor_shape.dimension_value(self._tensors[0].get_shape()[0]))
+    for t in self._tensors[1:]:
+      batch_dim.assert_is_compatible_with(
+          tensor_shape.Dimension(
+              tensor_shape.dimension_value(t.get_shape()[0])))
+
+    variant_tensor = gen_dataset_ops.tensor_slice_dataset(
+        self._tensors,
+        output_shapes=structure.get_flat_tensor_shapes(self._structure),
+        is_files=is_files,
+        metadata=self._metadata.SerializeToString())
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/from_tensors_op.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.from_tensors`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import structure
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _from_tensors(tensors, name):  # pylint: disable=unused-private-name
+  return _TensorDataset(tensors, name)
+
+
+class _TensorDataset(dataset_ops.DatasetSource):
+  """A `Dataset` with a single element."""
+
+  def __init__(self, element, name=None):
+    """See `tf.data.Dataset.from_tensors` for details."""
+    element = structure.normalize_element(element)
+    self._structure = structure.type_spec_from_value(element)
+    self._tensors = structure.to_tensor_list(self._structure, element)
+    self._name = name
+    variant_tensor = gen_dataset_ops.tensor_dataset(
+        self._tensors,
+        output_shapes=structure.get_flat_tensor_shapes(self._structure),
+        metadata=self._metadata.SerializeToString())
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/group_by_window_op.py

@@ -0,0 +1,132 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.group_by_window`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _group_by_window(input_dataset,  # pylint: disable=unused-private-name
+                     key_func,
+                     reduce_func,
+                     window_size=None,
+                     window_size_func=None,
+                     name=None):
+  """See `Dataset.group_by_window()` for details."""
+
+  if (window_size is not None and window_size_func or
+      not (window_size is not None or window_size_func)):
+    raise ValueError("Either the `window_size` argument or the "
+                     "`window_size_func` argument must be specified.")
+
+  if window_size is not None:
+
+    def constant_window_func(unused_key):
+      return ops.convert_to_tensor(window_size, dtype=dtypes.int64)
+
+    window_size_func = constant_window_func
+
+  assert window_size_func is not None
+
+  return _GroupByWindowDataset(
+      input_dataset, key_func, reduce_func, window_size_func, name=name)
+
+
+class _GroupByWindowDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that groups its input and performs a windowed reduction."""
+
+  def __init__(self,
+               input_dataset,
+               key_func,
+               reduce_func,
+               window_size_func,
+               name=None):
+    """See `group_by_window()` for details."""
+    self._input_dataset = input_dataset
+    self._make_key_func(key_func, input_dataset)
+    self._make_reduce_func(reduce_func, input_dataset)
+    self._make_window_size_func(window_size_func)
+    self._name = name
+    variant_tensor = ged_ops.group_by_window_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._key_func.function.captured_inputs,
+        self._reduce_func.function.captured_inputs,
+        self._window_size_func.function.captured_inputs,
+        key_func=self._key_func.function,
+        reduce_func=self._reduce_func.function,
+        window_size_func=self._window_size_func.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _make_window_size_func(self, window_size_func):
+    """Make wrapping defun for window_size_func."""
+
+    def window_size_func_wrapper(key):
+      return ops.convert_to_tensor(window_size_func(key), dtype=dtypes.int64)
+
+    self._window_size_func = structured_function.StructuredFunctionWrapper(
+        window_size_func_wrapper,
+        self._transformation_name(),
+        input_structure=tensor_spec.TensorSpec([], dtypes.int64))
+    if not self._window_size_func.output_structure.is_compatible_with(
+        tensor_spec.TensorSpec([], dtypes.int64)):
+      raise ValueError(f"Invalid `window_size_func`. `window_size_func` must "
+                       f"return a single `tf.int64` scalar tensor but its "
+                       f"return type is "
+                       f"{self._window_size_func.output_structure}.")
+
+  def _make_key_func(self, key_func, input_dataset):
+    """Make wrapping defun for key_func."""
+
+    def key_func_wrapper(*args):
+      return ops.convert_to_tensor(key_func(*args), dtype=dtypes.int64)
+
+    self._key_func = structured_function.StructuredFunctionWrapper(
+        key_func_wrapper, self._transformation_name(), dataset=input_dataset)
+    if not self._key_func.output_structure.is_compatible_with(
+        tensor_spec.TensorSpec([], dtypes.int64)):
+      raise ValueError(f"Invalid `key_func`. `key_func` must return a single "
+                       f"`tf.int64` scalar tensor but its return type is "
+                       f"{self._key_func.output_structure}.")
+
+  def _make_reduce_func(self, reduce_func, input_dataset):
+    """Make wrapping defun for reduce_func."""
+    nested_dataset = dataset_ops.DatasetSpec(input_dataset.element_spec)
+    input_structure = (tensor_spec.TensorSpec([], dtypes.int64), nested_dataset)
+    self._reduce_func = structured_function.StructuredFunctionWrapper(
+        reduce_func,
+        self._transformation_name(),
+        input_structure=input_structure)
+    if not isinstance(self._reduce_func.output_structure,
+                      dataset_ops.DatasetSpec):
+      raise TypeError(f"Invalid `reduce_func`. `reduce_func` must return a "
+                      f"single `tf.data.Dataset` object but its return type "
+                      f"is {self._reduce_func.output_structure}.")
+    # pylint: disable=protected-access
+    self._element_spec = (self._reduce_func.output_structure._element_spec)
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+  def _functions(self):
+    return [self._key_func, self._reduce_func, self._window_size_func]
+
+  def _transformation_name(self):
+    return "Dataset.group_by_window()"

tensorflow/python/data/ops/ignore_errors_op.py

@@ -0,0 +1,37 @@
+# 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 implementation of `tf.data.Dataset.ignore_errors`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops
+
+
+def _ignore_errors(input_dataset, log_warning=False, name=None):
+  return _IgnoreErrorsDataset(input_dataset, log_warning, name)
+
+
+class _IgnoreErrorsDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that drops erroneous elements from its input."""
+
+  def __init__(self, input_dataset, log_warning, name=None):
+    """See `Dataset.ignore_errors` for details."""
+    self._input_dataset = input_dataset
+    self._name = name
+    variant_tensor = (
+        gen_experimental_dataset_ops.ignore_errors_dataset(
+            self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+            log_warning=log_warning,
+            **self._flat_structure))
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/interleave_op.py

@@ -0,0 +1,170 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.interleave`."""
+
+import warnings
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _interleave(  # pylint: disable=unused-private-name
+    input_dataset,
+    map_func,
+    cycle_length=None,
+    block_length=None,
+    num_parallel_calls=None,
+    deterministic=None,
+    name=None):
+  """See `Dataset.interleave()` for details."""
+  if block_length is None:
+    block_length = 1
+
+  if cycle_length is None:
+    cycle_length = dataset_ops.AUTOTUNE
+
+  if num_parallel_calls is None or debug_mode.DEBUG_MODE:
+    if deterministic is not None and not debug_mode.DEBUG_MODE:
+      warnings.warn("The `deterministic` argument has no effect unless the "
+                    "`num_parallel_calls` argument is specified.")
+    return _InterleaveDataset(
+        input_dataset, map_func, cycle_length, block_length, name=name)
+  else:
+    return _ParallelInterleaveDataset(
+        input_dataset,
+        map_func,
+        cycle_length,
+        block_length,
+        num_parallel_calls,
+        deterministic=deterministic,
+        name=name)
+
+
+class _InterleaveDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that interleaves the result of transformed inputs."""
+
+  def __init__(self,
+               input_dataset,
+               map_func,
+               cycle_length,
+               block_length,
+               name=None):
+    """See `Dataset.interleave()` for details."""
+
+    self._input_dataset = input_dataset
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func, self._transformation_name(), dataset=input_dataset)
+    if not isinstance(self._map_func.output_structure, dataset_ops.DatasetSpec):
+      raise TypeError(
+          "The `map_func` argument must return a `Dataset` object. Got "
+          f"{dataset_ops.get_type(self._map_func.output_structure)!r}.")
+    self._structure = self._map_func.output_structure._element_spec  # pylint: disable=protected-access
+    self._cycle_length = ops.convert_to_tensor(
+        cycle_length, dtype=dtypes.int64, name="cycle_length")
+    self._block_length = ops.convert_to_tensor(
+        block_length, dtype=dtypes.int64, name="block_length")
+    self._name = name
+    variant_tensor = gen_dataset_ops.interleave_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,  # pylint: disable=protected-access
+        self._cycle_length,
+        self._block_length,
+        f=self._map_func.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+  def _transformation_name(self):
+    return "Dataset.interleave()"
+
+
+class _ParallelInterleaveDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over its input and interleaves the result.
+  """
+
+  def __init__(self,
+               input_dataset,
+               map_func,
+               cycle_length,
+               block_length,
+               num_parallel_calls,
+               buffer_output_elements=dataset_ops.AUTOTUNE,
+               prefetch_input_elements=dataset_ops.AUTOTUNE,
+               deterministic=None,
+               name=None):
+    """See `Dataset.interleave()` for details."""
+    self._input_dataset = input_dataset
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func, self._transformation_name(), dataset=input_dataset)
+    if not isinstance(self._map_func.output_structure, dataset_ops.DatasetSpec):
+      raise TypeError(
+          "The `map_func` argument must return a `Dataset` object. Got "
+          f"{dataset_ops.get_type(self._map_func.output_structure)!r}.")
+    self._structure = self._map_func.output_structure._element_spec  # pylint: disable=protected-access
+    self._cycle_length = ops.convert_to_tensor(
+        cycle_length, dtype=dtypes.int64, name="cycle_length")
+    self._block_length = ops.convert_to_tensor(
+        block_length, dtype=dtypes.int64, name="block_length")
+    self._buffer_output_elements = ops.convert_to_tensor(
+        buffer_output_elements,
+        dtype=dtypes.int64,
+        name="buffer_output_elements")
+    self._prefetch_input_elements = ops.convert_to_tensor(
+        prefetch_input_elements,
+        dtype=dtypes.int64,
+        name="prefetch_input_elements")
+
+    self._num_parallel_calls = ops.convert_to_tensor(
+        num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
+    if deterministic is None:
+      deterministic_string = "default"
+    elif deterministic:
+      deterministic_string = "true"
+    else:
+      deterministic_string = "false"
+
+    self._name = name
+    variant_tensor = gen_dataset_ops.parallel_interleave_dataset_v4(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,  # pylint: disable=protected-access
+        self._cycle_length,
+        self._block_length,
+        self._buffer_output_elements,
+        self._prefetch_input_elements,
+        self._num_parallel_calls,
+        f=self._map_func.function,
+        deterministic=deterministic_string,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+  def _transformation_name(self):
+    return "Dataset.interleave()"

tensorflow/python/data/ops/iterator_autograph.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+"""Autograph specific overrides for tf.data.ops."""
+import functools
+
+import numpy as np
+
+from tensorflow.python.autograph.operators import control_flow
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.framework import tensor_conversion
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import cond
+from tensorflow.python.util import nest
+
+
+# TODO(mdan): These checks should be easier. Fix the nest API.
+def _verify_spec_compatible(input_name, spec_name, input_, spec):
+  """Verifies that a symbol has a type compatible vith a given spec.
+
+  Here, compatibility is viewed in the general TensorFlow sense: that the dtypes
+  are the same after implicit conversion, if both are tensors.
+
+  This verifier ensures consistent treatment of types across AutoGraph.
+
+  Args:
+    input_name: A name to use for `input_` in error messages.
+    spec_name: A name to use for `spec` in error messages.
+    input_: Any, value to verify.
+    spec: TypeSpec that `input_` must be compatible with.
+
+  Raises:
+    ValueError if the two types have been determined not to be compatible.
+  """
+  assert isinstance(spec, tensor_spec.TensorSpec)
+  if input is None:
+    # TODO(mdan): raise from None when switching to Py3.
+    raise ValueError("{} cannot be None".format(input_name))
+
+  # TODO(mdan): Use TensorCompatible when ready.
+  if isinstance(input_, (bool, int, float, str, np.ndarray)):
+    input_ = tensor_conversion.convert_to_tensor_v2(input_)
+
+  input_dtype = getattr(input_, "dtype", None)
+
+  if input_dtype != spec.dtype:
+    input_dtype_str = "no dtype" if input_dtype is None else str(input_dtype)
+
+    raise TypeError(
+        "{} must have the same dtype as {}. Expected {}, got {}".format(
+            input_name, spec_name, spec.dtype, input_dtype_str
+        )
+    )
+
+
+def _verify_structure_compatible(input_name, spec_name, input_, spec):
+  """Verifies that possibly-structured symbol has types compatible vith another.
+
+  See _verify_spec_compatible for a more concrete meaning of "compatible".
+  Unspec _verify_spec_compatible, which handles singular Tensor-spec objects,
+  verify_structures_compatible can process structures recognized by tf.nest.
+
+  Args:
+    input_name: A name to use for `input_` in error messages.
+    spec_name: A name to use for `spec` in error messages.
+    input_: Any, value to verify. May, but doesn't need to, be a structure.
+    spec: Any, value that `input_` must be compatible with. May, but doesn't
+      need to, be a structure.
+
+  Raises:
+    ValueError if the two types have been determined not to be compatible.
+  """
+  try:
+    nest.assert_same_structure(input_, spec, expand_composites=True)
+  except (ValueError, TypeError) as e:
+    raise TypeError(
+        "{} must have the same element structure as {}.\n\n{}".format(
+            input_name, spec_name, str(e)
+        )
+    ) from e
+
+  nest.map_structure(
+      functools.partial(_verify_spec_compatible, input_name, spec_name), input_,
+      spec)
+
+
+def _next_tf_iterator(iterator, default=py_builtins.UNSPECIFIED):
+  if default is py_builtins.UNSPECIFIED:
+    # Without a default, fall back to the "normal" behavior which raises
+    # a runtime exception.
+    return next(iterator)
+  opt_iterate = iterator.get_next_as_optional()
+  _verify_structure_compatible(
+      "the default argument", "the iterate", default, iterator.element_spec
+  )
+  return cond.cond(
+      opt_iterate.has_value(), opt_iterate.get_value, lambda: default
+  )
+
+
+def register_overrides():
+  py_builtins.next_registry.register(
+      iterator_ops.OwnedIterator, _next_tf_iterator
+  )
+  control_flow.for_loop_registry.register(
+      iterator_ops.OwnedIterator, control_flow._tf_iterator_for_stmt  # pylint: disable=protected-access
+  )

tensorflow/python/data/ops/load_op.py

@@ -0,0 +1,279 @@
+# 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.
+# ==============================================================================
+"""Implementation of LoadDataset in Python."""
+import multiprocessing
+import os
+import time
+from typing import Any, Callable, Optional, Union
+
+from absl import logging
+
+from google.protobuf import message
+from google.protobuf import text_format
+from tensorflow.core.protobuf import snapshot_pb2
+from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.platform import gfile
+# TODO(b/238903802): Use TypeSpec serialization methods directly.
+from tensorflow.python.saved_model import nested_structure_coder
+
+# For distributed snapshot load V2, retries loading after this time, if the
+# snapshot is not ready yet.
+_RETRY_INTERVAL_SEC = 5
+
+
+def _load(  # pylint: disable=unused-private-name
+    path: str,
+    element_spec: Any,
+    compression: Optional[str],
+    reader_func: Optional[Callable[[dataset_ops.Dataset], dataset_ops.Dataset]],
+    wait: bool,
+) -> dataset_ops.Dataset:
+  """Loads dataset from tf.data snapshot."""
+
+  if wait:
+    return _load_with_retry(path, element_spec, compression, reader_func)
+
+  if reader_func is None:
+    reader_func = lambda datasets: datasets.interleave(  # pylint:disable=g-long-lambda
+        lambda x: x,
+        cycle_length=multiprocessing.cpu_count(),
+        num_parallel_calls=dataset_ops.AUTOTUNE)
+
+  distributed_snapshot_metadata = _load_distributed_snapshot_metadata(path)
+  if distributed_snapshot_metadata:
+    _validate_snapshot(
+        path, distributed_snapshot_metadata, element_spec, compression)
+    return _load_distributed_snapshot(
+        path, distributed_snapshot_metadata, reader_func)
+
+  if element_spec is None:
+    element_spec = _load_element_spec(path)
+  return _LoadDataset(path, element_spec, compression, reader_func)
+
+
+def _load_with_retry(  # pylint: disable=unused-private-name
+    path: str,
+    element_spec: Any = None,
+    compression: Optional[str] = None,
+    reader_func: Optional[
+        Callable[[dataset_ops.Dataset], dataset_ops.Dataset]] = None,
+) -> dataset_ops.Dataset:
+  """Tries loading the snapshot. Retries if not found."""
+
+  while True:
+    try:
+      dataset = dataset_ops.Dataset.load(
+          path=path,
+          element_spec=element_spec,
+          compression=compression,
+          reader_func=reader_func,
+          wait=False)
+      logging.info("Load tf.data snapshot at %s.", path)
+      return dataset
+    except (errors.NotFoundError, FileNotFoundError):
+      logging.info(
+          "Could not find tf.data snapshot at %s. Will wait and retry.", path)
+      time.sleep(_RETRY_INTERVAL_SEC)
+
+
+def _load_distributed_snapshot_metadata(
+    path: str,
+) -> Optional[snapshot_pb2.DistributedSnapshotMetadata]:
+  """Reads the distributed snapshot metadata.
+
+  Args:
+    path: Base path of the snapshot.
+
+  Returns:
+    DistributedSnapshotMetadata if the snapshot is a distributed snapshot.
+    Returns None if it is a non-distributed snapshot.
+  """
+  metadata_file = _pywrap_snapshot_utils.TF_DATA_SnapshotMetadataFilePath(path)
+  if not gfile.Exists(metadata_file):
+    return None
+
+  try:
+    with gfile.GFile(metadata_file, "r") as f:
+      return text_format.ParseLines(
+          f, snapshot_pb2.DistributedSnapshotMetadata())
+  except (
+      errors.NotFoundError,
+      text_format.ParseError,
+      message.DecodeError,
+      UnicodeDecodeError):
+    return None
+
+
+def _load_distributed_snapshot(
+    path: str,
+    metadata: snapshot_pb2.DistributedSnapshotMetadata,
+    reader_func: Callable[[dataset_ops.Dataset], dataset_ops.Dataset],
+) -> dataset_ops.Dataset:
+  """Loads a distributed snapshot."""
+
+  dataset = _ListSnapshotChunksDataset(path)
+  dataset = dataset.map(
+      lambda chunk_file: _SnapshotChunkDataset(  # pylint:disable=g-long-lambda
+          chunk_file,
+          element_spec=_parse_element_spec(metadata.element_spec),
+          compression=metadata.compression))
+  return reader_func(dataset)
+
+
+def _load_element_spec(path: str) -> Any:
+  """Loads the dataset element spec.
+
+  Args:
+    path: Base path of the snapshot.
+
+  Returns:
+    Dataset element_spec.
+
+  Raises:
+    NotFoundError if the element spec file does not exist or cannot be decoded.
+  """
+  dataset_spec_filename = os.path.join(path, dataset_ops.DATASET_SPEC_FILENAME)
+  if not gfile.Exists(dataset_spec_filename):
+    raise errors.NotFoundError(
+        node_def=None, op=None,
+        message="tf.data snapshot element_spec file not found: "
+                f"{dataset_spec_filename}.")
+
+  with gfile.GFile(dataset_spec_filename, "rb") as f:
+    encoded_spec = f.read()
+  try:
+    return _parse_element_spec(encoded_spec)
+  except nested_structure_coder.NotEncodableError as e:
+    raise errors.NotFoundError(
+        node_def=None, op=None,
+        message="tf.data snapshot element_spec file not found or invalid: "
+                f"{dataset_spec_filename}.") from e
+
+
+def _parse_element_spec(encoded_element_spec: Union[bytes, str]) -> Any:
+  struct_pb = nested_structure_coder.struct_pb2.StructuredValue()
+  struct_pb.ParseFromString(encoded_element_spec)
+  return nested_structure_coder.decode_proto(struct_pb)
+
+
+class _LoadDataset(dataset_ops.DatasetSource):
+  """A dataset that loads previously saved dataset."""
+
+  def __init__(
+      self,
+      path: str,
+      element_spec: Any,
+      compression: str,
+      reader_func: Callable[[dataset_ops.Dataset], dataset_ops.Dataset]):
+    self._path = path
+    self._element_spec = element_spec
+    self._compression = compression
+    self._reader_func = structured_function.StructuredFunctionWrapper(
+        reader_func,
+        "load()",
+        # Dataset of datasets of input elements
+        input_structure=dataset_ops.DatasetSpec(
+            dataset_ops.DatasetSpec(self._element_spec)))
+
+    variant_tensor = ged_ops.load_dataset(
+        path,
+        reader_func_other_args=self._reader_func.function.captured_inputs,
+        compression=compression,
+        reader_func=self._reader_func.function,
+        **self._flat_structure)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self) -> Any:
+    return self._element_spec
+
+
+class _SnapshotChunkDataset(dataset_ops.DatasetSource):
+  """A dataset for one chunk file from a tf.data distributed snapshot."""
+
+  def __init__(self, chunk_file: str, element_spec: Any, compression: str):
+    self._chunk_file = chunk_file
+    self._element_spec = element_spec
+    variant_tensor = ged_ops.snapshot_chunk_dataset(
+        chunk_file,
+        compression=compression,
+        **self._flat_structure)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self) -> Any:
+    return self._element_spec
+
+
+class _ListSnapshotChunksDataset(dataset_ops.DatasetSource):
+  """A dataset for listing snapshot chunk files.
+
+  It supports listing partially written snapshots. When a snapshot is being
+  written, it returns the currently available chunk files.
+  """
+
+  def __init__(self, snapshot_path: str):
+    self._snapshot_path = snapshot_path
+    variant_tensor = ged_ops.list_snapshot_chunks_dataset(
+        snapshot_path, **self._flat_structure)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self) -> tensor_spec.TensorSpec:
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+def _validate_snapshot(
+    path: str,
+    metadata: snapshot_pb2.DistributedSnapshotMetadata,
+    element_spec: Any,
+    compression: str) -> None:
+  """Validates a tf.data distributed snapshot.
+
+  Args:
+    path: Root path of the distributed snapshot.
+    metadata: The DistributedSnapshotMetadata of the snapshot.
+    element_spec: Dataset element_spec.
+    compression: Compression method used for saving.
+
+  Raises:
+    ValueError if the snapshot is invalid.
+  """
+
+  error_file = _pywrap_snapshot_utils.TF_DATA_SnapshotErrorFilePath(path)
+  if gfile.Exists(error_file):
+    with gfile.GFile(error_file, "r") as f:
+      raise ValueError(
+          f"Failed to load tf.data snapshot at {path}. The save job failed to "
+          f"write it. Status: {f.read()}")
+
+  snapshot_element_spec = _parse_element_spec(metadata.element_spec)
+  if element_spec and element_spec != snapshot_element_spec:
+    raise ValueError(
+        f"Failed to load tf.data snapshot at {path}. User specified "
+        f"element_spec {element_spec}, but the actual element_spec is "
+        f"{snapshot_element_spec}.")
+
+  if compression and compression != metadata.compression:
+    raise ValueError(
+        f"Failed to load tf.data snapshot at {path}. User specified "
+        f"compression {compression}, but the actual compression is "
+        f"{metadata.compression}.")

tensorflow/python/data/ops/map_op.py

@@ -0,0 +1,238 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.map`."""
+
+import warnings
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _map_v2(
+    input_dataset,  # pylint: disable=unused-private-name
+    map_func,
+    num_parallel_calls=None,
+    deterministic=None,
+    synchronous=None,
+    use_unbounded_threadpool=None,
+    name=None,
+):
+  """See `Dataset.map()` for details."""
+  if num_parallel_calls is None or debug_mode.DEBUG_MODE:
+    if deterministic is not None and not debug_mode.DEBUG_MODE:
+      warnings.warn(
+          "The `deterministic` argument has no effect unless the "
+          "`num_parallel_calls` argument is specified."
+      )
+    return _MapDataset(
+        input_dataset,
+        map_func,
+        preserve_cardinality=True,
+        force_synchronous=False if synchronous is None else synchronous,
+        name=name,
+    )
+  else:
+    if synchronous:
+      raise ValueError(
+          "`synchronous` is not supported with `num_parallel_calls`, but"
+          " `num_parallel_calls` was set to ",
+          num_parallel_calls,
+      )
+    return _ParallelMapDataset(
+        input_dataset,
+        map_func,
+        num_parallel_calls=num_parallel_calls,
+        deterministic=deterministic,
+        preserve_cardinality=True,
+        use_unbounded_threadpool=use_unbounded_threadpool,
+        name=name)
+
+
+def _map_v1(
+    input_dataset,  # pylint: disable=unused-private-name
+    map_func,
+    num_parallel_calls=None,
+    deterministic=None,
+    synchronous=None,
+    use_unbounded_threadpool=None,  # pylint: disable=unused-argument
+):
+  """See `Dataset.map()` for details."""
+  if num_parallel_calls is None or debug_mode.DEBUG_MODE:
+    return dataset_ops.DatasetV1Adapter(
+        _MapDataset(
+            input_dataset,
+            map_func,
+            preserve_cardinality=False,
+            force_synchronous=False if synchronous is None else synchronous,
+        )
+    )
+  else:
+    if synchronous:
+      raise ValueError(
+          "`synchronous` is not supported with `num_parallel_calls`, but"
+          " `num_parallel_calls` was set to ",
+          num_parallel_calls,
+      )
+    return dataset_ops.DatasetV1Adapter(
+        _ParallelMapDataset(
+            input_dataset,
+            map_func,
+            num_parallel_calls,
+            deterministic,
+            preserve_cardinality=False,
+            use_unbounded_threadpool=False))
+
+
+def _map_v1_with_legacy_function(  # pylint: disable=unused-private-name
+    input_dataset,
+    map_func,
+    num_parallel_calls=None,
+    deterministic=None,
+    synchronous=False,
+):
+  """See `Dataset.map()` for details."""
+  if num_parallel_calls is None:
+    if deterministic is not None:
+      warnings.warn("The `deterministic` argument has no effect unless the "
+                    "`num_parallel_calls` argument is specified.")
+    return dataset_ops.DatasetV1Adapter(
+        _MapDataset(
+            input_dataset,
+            map_func,
+            force_synchronous=synchronous,
+            preserve_cardinality=False,
+            use_legacy_function=True,
+        )
+    )
+  else:
+    if synchronous:
+      raise ValueError(
+          "`synchronous` is not supported with `num_parallel_calls`, but"
+          " `num_parallel_calls` was set to ",
+          num_parallel_calls,
+      )
+    return dataset_ops.DatasetV1Adapter(
+        _ParallelMapDataset(
+            input_dataset,
+            map_func,
+            num_parallel_calls,
+            deterministic,
+            preserve_cardinality=False,
+            use_legacy_function=True,
+            use_unbounded_threadpool=False))
+
+
+class _MapDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over elements in its input."""
+
+  def __init__(
+      self,
+      input_dataset,
+      map_func,
+      force_synchronous=False,
+      use_inter_op_parallelism=True,
+      preserve_cardinality=True,
+      use_legacy_function=False,
+      name=None,
+  ):
+    self._input_dataset = input_dataset
+    self._use_inter_op_parallelism = use_inter_op_parallelism
+    self._preserve_cardinality = preserve_cardinality
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func,
+        self._transformation_name(),
+        dataset=input_dataset,
+        use_legacy_function=use_legacy_function)
+    self._force_synchronous = force_synchronous
+    self._name = name
+    variant_tensor = gen_dataset_ops.map_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        f=self._map_func.function,
+        use_inter_op_parallelism=self._use_inter_op_parallelism,
+        preserve_cardinality=self._preserve_cardinality,
+        force_synchronous=self._force_synchronous,
+        **self._common_args
+    )
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._map_func.output_structure
+
+  def _transformation_name(self):
+    return "Dataset.map()"
+
+
+class _ParallelMapDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over elements in its input in parallel."""
+
+  def __init__(self,
+               input_dataset,
+               map_func,
+               num_parallel_calls,
+               deterministic,
+               use_inter_op_parallelism=True,
+               preserve_cardinality=False,
+               use_legacy_function=False,
+               use_unbounded_threadpool=False,
+               name=None):
+    """See `Dataset.map()` for details."""
+    self._input_dataset = input_dataset
+    self._use_inter_op_parallelism = use_inter_op_parallelism
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func,
+        self._transformation_name(),
+        dataset=input_dataset,
+        use_legacy_function=use_legacy_function)
+    if deterministic is None:
+      self._deterministic = "default"
+    elif deterministic:
+      self._deterministic = "true"
+    else:
+      self._deterministic = "false"
+    self._preserve_cardinality = preserve_cardinality
+    self._num_parallel_calls = ops.convert_to_tensor(
+        num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
+    self._use_unbounded_threadpool = use_unbounded_threadpool
+    self._name = name
+    variant_tensor = gen_dataset_ops.parallel_map_dataset_v2(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        f=self._map_func.function,
+        num_parallel_calls=self._num_parallel_calls,
+        deterministic=self._deterministic,
+        use_inter_op_parallelism=self._use_inter_op_parallelism,
+        preserve_cardinality=self._preserve_cardinality,
+        use_unbounded_threadpool=self._use_unbounded_threadpool,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._map_func.output_structure
+
+  def _transformation_name(self):
+    return "Dataset.map()"

tensorflow/python/data/ops/multi_device_iterator_ops.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.
+# ==============================================================================
+"""Python wrapper for prefetching_ops."""
+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.ops import prefetch_op
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import type_spec
+from tensorflow.python.framework import type_utils
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import resource_variable_ops
+
+
+class _PerDeviceGenerator(dataset_ops.DatasetV2):
+  """A `dummy` generator dataset."""
+
+  def __init__(self, shard_num, multi_device_iterator_resource, incarnation_id,
+               source_device, element_spec, iterator_is_anonymous):
+    self._element_spec = element_spec
+    self._name = f"device_generator_{shard_num}"
+
+    multi_device_iterator_string_handle = (
+        gen_dataset_ops.multi_device_iterator_to_string_handle(
+            multi_device_iterator_resource))
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(autograph=False)  # Pure graph code.
+    def _init_func():
+      return multi_device_iterator_string_handle
+
+    init_func_concrete = _init_func.get_concrete_function()
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(autograph=False)  # Pure graph code.
+    def _remote_init_func():
+      return functional_ops.remote_call(
+          target=source_device,
+          args=init_func_concrete.captured_inputs,
+          Tout=[dtypes.string],
+          f=init_func_concrete)
+
+    self._init_func = _remote_init_func.get_concrete_function()
+    self._init_captured_args = self._init_func.captured_inputs
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.string)],
+        autograph=False)  # Pure graph code.
+    def _next_func(string_handle):
+      # pylint: disable=protected-access
+      multi_device_iterator = (
+          gen_dataset_ops.multi_device_iterator_from_string_handle(
+              string_handle=string_handle,
+              output_types=structure.get_flat_tensor_types(self._element_spec),
+              output_shapes=structure.get_flat_tensor_shapes(
+                  self._element_spec)))
+      return gen_dataset_ops.multi_device_iterator_get_next_from_shard(
+          multi_device_iterator=multi_device_iterator,
+          shard_num=shard_num,
+          incarnation_id=incarnation_id,
+          output_types=structure.get_flat_tensor_types(self._element_spec),
+          output_shapes=structure.get_flat_tensor_shapes(self._element_spec))
+
+    next_func_concrete = _next_func.get_concrete_function()
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.string)],
+        experimental_attributes={"experimental_ints_on_device": True},
+        autograph=False)  # Pure graph code.
+    def _remote_next_func(string_handle):
+      return_values = functional_ops.remote_call(
+          target=source_device,
+          args=[string_handle] + next_func_concrete.captured_inputs,
+          Tout=structure.get_flat_tensor_types(self._element_spec),
+          f=next_func_concrete)
+      # Add full type information to the graph so that the RemoteCall op
+      # can determine for each of its outputs whether or not they are ragged
+      # tensors (or other types that use variants) that contain strings
+      # (or other host memory types). Then RemoteCall can
+      # appropriately set AllocatorAttributes to control copies so
+      # strings/host memory types stay on CPU.
+      fulltype_list = type_utils.fulltypes_for_flat_tensors(self._element_spec)
+      fulltype = type_utils.fulltype_list_to_product(fulltype_list)
+      for return_value in return_values:
+        return_value.op.experimental_set_type(fulltype)
+      return return_values
+
+    self._next_func = _remote_next_func.get_concrete_function()
+    self._next_captured_args = self._next_func.captured_inputs
+
+    if iterator_is_anonymous:
+      self._next_captured_args = self._next_captured_args + [
+          multi_device_iterator_resource
+      ]
+
+    self._incarnation_id_index = -1
+    for i, arg in enumerate(self._next_captured_args):
+      if arg is incarnation_id:
+        self._incarnation_id_index = i
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.string)],
+        autograph=False)  # Pure graph code.
+    def _finalize_func(unused_string_handle):
+      return array_ops.constant(0, dtypes.int64)
+
+    finalize_func_concrete = _finalize_func.get_concrete_function()
+
+    # TODO(b/124254153): Enable autograph once the overhead is low enough.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.string)],
+        autograph=False)  # Pure graph code.
+    def _remote_finalize_func(string_handle):
+      return functional_ops.remote_call(
+          target=source_device,
+          args=[string_handle] + finalize_func_concrete.captured_inputs,
+          Tout=[dtypes.int64],
+          f=finalize_func_concrete)
+
+    self._finalize_func = _remote_finalize_func.get_concrete_function()
+    self._finalize_captured_args = self._finalize_func.captured_inputs
+
+    variant_tensor = gen_dataset_ops.generator_dataset(
+        self._init_captured_args,
+        self._next_captured_args,
+        self._finalize_captured_args,
+        init_func=self._init_func,
+        next_func=self._next_func,
+        finalize_func=self._finalize_func,
+        **self._flat_structure)
+    super(_PerDeviceGenerator, self).__init__(variant_tensor)
+
+  def _inputs(self):
+    # TODO(b/116506223): Determine which datasets should be used as inputs here.
+    return []
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+class _ReincarnatedPerDeviceGenerator(dataset_ops.DatasetV2):
+  """Creates a _PerDeviceGenerator-like dataset with a new incarnation_id.
+
+  Re-uses the functions from the provided per_device_dataset and just switches
+  out the function argument corresponding to the incarnation_id.
+  """
+
+  def __init__(self, per_device_dataset, incarnation_id):
+    # pylint: disable=protected-access
+    if hasattr(per_device_dataset, "_name"):
+      self._name = per_device_dataset._name
+    self._element_spec = per_device_dataset.element_spec
+    self._init_func = per_device_dataset._init_func
+    self._init_captured_args = self._init_func.captured_inputs
+
+    self._next_func = per_device_dataset._next_func
+    self._next_captured_args = per_device_dataset._next_captured_args
+    # The captured arguments to the next_func are string_handle, incarnation_id.
+    # We update the incarnation id to the new one.
+    self._next_captured_args[
+        per_device_dataset._incarnation_id_index] = incarnation_id
+
+    self._finalize_func = per_device_dataset._finalize_func
+    self._finalize_captured_args = per_device_dataset._finalize_captured_args
+
+    variant_tensor = gen_dataset_ops.generator_dataset(
+        self._init_captured_args,
+        self._next_captured_args,
+        self._finalize_captured_args,
+        init_func=self._init_func,
+        next_func=self._next_func,
+        finalize_func=self._finalize_func,
+        **self._flat_structure)
+    super(_ReincarnatedPerDeviceGenerator, self).__init__(variant_tensor)
+
+  def _inputs(self):
+    # TODO(b/116506223): Determine which datasets should be used as inputs here.
+    return []
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+def _create_device_dataset(prototype_ds, incarnation_id, prefetch_buffer_size,
+                           experimental_slack):
+  """Uses _prototype_device_datasets[i] to build a dataset for the device."""
+  ds = _ReincarnatedPerDeviceGenerator(prototype_ds, incarnation_id)
+  if prefetch_buffer_size > 0:
+    if experimental_slack:
+      ds = prefetch_op._PrefetchDataset(  # pylint: disable=protected-access
+          ds, prefetch_buffer_size, slack_period=1)
+    else:
+      ds = ds.prefetch(prefetch_buffer_size, name="device_prefetch")
+  return ds
+
+
+class MultiDeviceIterator:
+  """An iterator over multiple devices."""
+
+  def __init__(self,
+               dataset,
+               devices,
+               max_buffer_size=1,
+               prefetch_buffer_size=1,
+               source_device="/cpu:0"):
+    """Constructs a MultiDeviceIterator.
+
+    Args:
+      dataset: The input dataset to be iterated over.
+      devices: The list of devices to fetch data to.
+      max_buffer_size: Maximum size of the host side per device buffer to keep.
+      prefetch_buffer_size: if > 0, then we setup a buffer on each device to
+        prefetch into.
+      source_device: The host device to place the `dataset` on.  In order to
+        prevent deadlocks, if the prefetch_buffer_size is greater than the
+        max_buffer_size, we set the max_buffer_size to prefetch_buffer_size.
+    """
+    options = options_lib.Options()
+    options.experimental_distribute.num_devices = len(devices)
+    # If `prefetch_buffer_size` is 0, we turn off the `inject_prefetch`
+    # optimization to prevent potentially introducing asynchrony.
+    if prefetch_buffer_size == 0:
+      options.experimental_optimization.inject_prefetch = False
+    dataset = dataset.with_options(options)
+    self._dataset = dataset._apply_debug_options()  # pylint: disable=protected-access
+    self._experimental_slack = dataset.options().experimental_slack
+    self._devices = devices
+    self._source_device = source_device
+    self._source_device_tensor = ops.convert_to_tensor(source_device)
+    self._max_buffer_size = max_buffer_size
+    self._prefetch_buffer_size = prefetch_buffer_size
+
+    if self._prefetch_buffer_size > self._max_buffer_size:
+      self._max_buffer_size = self._prefetch_buffer_size
+
+    # Create the MultiDeviceIterator.
+    with ops.device(self._source_device):
+      # TODO(b/121378567): Get rid of this shared_name hack.
+      shared_name = ""
+      if context.executing_eagerly():
+        shared_name = context.anonymous_name()
+      self._multi_device_iterator_resource = (
+          gen_dataset_ops.multi_device_iterator(
+              devices=self._devices,
+              shared_name=shared_name,
+              container="",
+              **self._dataset._flat_structure))  # pylint: disable=protected-access
+      if context.executing_eagerly():
+        # Delete the resource when this object is deleted
+        self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
+            handle=self._multi_device_iterator_resource,
+            handle_device=self._source_device)
+
+      # The incarnation ID is used to ensure consistency between the per-device
+      # iterators and the multi-device iterator.
+      self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
+          self._dataset._variant_tensor,  # pylint: disable=protected-access
+          self._multi_device_iterator_resource,
+          max_buffer_size=self._max_buffer_size)
+
+    self._prototype_device_datasets = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        ds = _PerDeviceGenerator(
+            i,
+            self._multi_device_iterator_resource,
+            self._incarnation_id,
+            self._source_device_tensor,
+            self._dataset.element_spec,
+            iterator_is_anonymous=False)
+        self._prototype_device_datasets.append(ds)
+
+    # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to
+    # initialize the device side of the pipeline. This would allow the
+    # MultiDeviceIterator to choose, for example, to move some transformations
+    # into the device side from its input. It might be useful in rewriting.
+    # Create the per device iterators.
+    self._device_iterators = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        ds = _create_device_dataset(self._prototype_device_datasets[i],
+                                    self._incarnation_id,
+                                    self._prefetch_buffer_size,
+                                    self._experimental_slack)
+        if context.executing_eagerly():
+          self._device_iterators.append(dataset_ops.make_one_shot_iterator(ds))
+        else:
+          self._device_iterators.append(
+              dataset_ops.make_initializable_iterator(ds))
+
+    if not context.executing_eagerly():
+      device_iterator_initializers = [
+          iterator.initializer for iterator in self._device_iterators
+      ]
+      self._initializer = control_flow_ops.group(*device_iterator_initializers)
+
+  def get_next(self, device=None):
+    """Returns the next element given a `device`, else returns all in a list."""
+    if device is not None:
+      index = self._devices.index(device)
+      return self._device_iterators[index].get_next()
+
+    result = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        result.append(self._device_iterators[i].get_next())
+    return result
+
+  def get_next_as_optional(self):
+    result = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        result.append(self._device_iterators[i].get_next_as_optional())
+    return result
+
+  @property
+  def initializer(self):
+    if context.executing_eagerly():
+      return control_flow_ops.no_op()
+    return self._initializer
+
+  def _eager_reset(self):
+    """Resets the MultiDeviceIterator in eager mode."""
+    if not ops.executing_eagerly_outside_functions():
+      raise ValueError(
+          "Resetting a multi-device iterator is only supported in the eager "
+          "mode.")
+    # pylint: disable=protected-access
+    self._incarnation_id = gen_dataset_ops.multi_device_iterator_init(
+        self._dataset._variant_tensor,
+        self._multi_device_iterator_resource,
+        max_buffer_size=self._max_buffer_size)
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        ds = _create_device_dataset(self._prototype_device_datasets[i],
+                                    self._incarnation_id,
+                                    self._prefetch_buffer_size,
+                                    self._experimental_slack)
+        # Reset the device iterator resources with the new dataset.
+        ds_variant = ds._variant_tensor
+        gen_dataset_ops.make_iterator(
+            ds_variant, self._device_iterators[i]._iterator_resource)
+
+  @property
+  def element_spec(self):
+    return self._dataset.element_spec
+
+
+class MultiDeviceIteratorSpec(type_spec.TypeSpec):
+  """Type specification for `OwnedMultiDeviceIterator`."""
+
+  __slots__ = ["_devices", "_source_device", "_element_spec"]
+
+  def __init__(self, devices, source_device, element_spec):
+    self._devices = devices
+    self._source_device = source_device
+    self._element_spec = element_spec
+
+  @property
+  def value_type(self):
+    return OwnedMultiDeviceIterator
+
+  def _serialize(self):
+    return (tuple(self._devices), self._source_device, self._element_spec)
+
+  @property
+  def _component_specs(self):
+    specs = [
+        tensor_spec.TensorSpec([], dtypes.resource),
+    ]
+    for _ in range(len(self._devices)):
+      specs.append(iterator_ops.IteratorSpec(self._element_spec))
+    return specs
+
+  def _to_components(self, value):
+    # pylint: disable=protected-access
+    c = [value._multi_device_iterator_resource]
+    c.extend(value._device_iterators)
+    return c
+
+  def _from_components(self, components):
+    return OwnedMultiDeviceIterator(
+        dataset=None,
+        devices=self._devices,
+        source_device=self._source_device,
+        components=components,
+        element_spec=self._element_spec)
+
+  @staticmethod
+  def from_value(value):
+    # pylint: disable=protected-access
+    return MultiDeviceIteratorSpec(
+        value._devices,
+        value._source_device,
+        value.element_spec)
+
+
+class OwnedMultiDeviceIterator(composite_tensor.CompositeTensor):
+  """An iterator over multiple devices.
+
+  The multi-device iterator resource created through `OwnedMultiDeviceIterator`
+  is owned by the Python object and the life time of the underlying resource is
+  tied to the life time of the `OwnedMultiDeviceIterator` object. This makes
+  `OwnedMultiDeviceIterator` appropriate for use in eager mode and inside of
+  tf.functions.
+  """
+
+  def __init__(self,
+               dataset=None,
+               devices=None,
+               max_buffer_size=1,
+               prefetch_buffer_size=1,
+               source_device="/cpu:0",
+               components=None,
+               element_spec=None):
+    """Constructs an owned MultiDeviceIterator object.
+
+    Args:
+      dataset: The input dataset to be iterated over.
+      devices: (Required.) The list of devices to fetch data to.
+      max_buffer_size: Maximum size of the host side per device buffer to keep.
+      prefetch_buffer_size: if > 0, then we setup a buffer on each device to
+        prefetch into.
+      source_device: The host device to place the `dataset` on.  In order to
+        prevent deadlocks, if the prefetch_buffer_size is greater than the
+        max_buffer_size, we set the max_buffer_size to prefetch_buffer_size.
+      components: Tensor components to construct the MultiDeviceIterator from.
+      element_spec: A (nested) structure of `tf.TypeSpec` objects that
+        represents the type specification of elements of the iterator.
+
+    Raises:
+      RuntimeError: If executed in graph mode or outside of function building
+        mode.
+      ValueError: If any of the following happens:
+        - `devices` is `None`
+        - `dataset` is `None` and either `components` or `element_spec` is
+          `None`
+        - `dataset` is not None and either `components` or `element_spec` is
+          provided
+    """
+    if not context.executing_eagerly() and not ops.inside_function():
+      raise RuntimeError("OwnedMultiDeviceIterator is only supported inside of "
+                         "tf.function or when eager execution is enabled.")
+    if devices is None:
+      raise ValueError("`devices` must be provided.")
+
+    if dataset is None:
+      if (components is None or element_spec is None):
+        raise ValueError(
+            "When `dataset` is not provided, both `components` and "
+            "`element_spec` must be specified.")
+      self._element_spec = element_spec
+      self._devices = devices
+      self._source_device = source_device
+      self._multi_device_iterator_resource = components[0]
+      self._device_iterators = components[1:]
+    else:
+      if (components is not None or element_spec is not None):
+        raise ValueError(
+            "When `dataset` is provided, `element_spec` and `components` must "
+            "not be specified.")
+      options = options_lib.Options()
+      options.experimental_distribute.num_devices = len(devices)
+      # If `prefetch_buffer_size` is 0, we turn off the `inject_prefetch`
+      # optimization to prevent potentially introducing asynchrony.
+      if prefetch_buffer_size == 0:
+        options.experimental_optimization.inject_prefetch = False
+      dataset = dataset.with_options(options)
+      dataset = dataset._apply_debug_options()  # pylint: disable=protected-access
+      self._element_spec = dataset.element_spec
+      experimental_slack = dataset.options().experimental_slack
+      self._devices = devices
+      self._source_device = source_device
+      source_device_tensor = ops.convert_to_tensor(self._source_device)
+
+      if prefetch_buffer_size > max_buffer_size:
+        max_buffer_size = prefetch_buffer_size
+
+      # Create the MultiDeviceIterator.
+      with ops.device(self._source_device):
+        self._multi_device_iterator_resource = (
+            gen_dataset_ops.anonymous_multi_device_iterator_v3(
+                devices=self._devices, **dataset._flat_structure))  # pylint: disable=protected-access
+
+        # The incarnation ID is used to ensure consistency between the
+        # per-device iterators and the multi-device iterator.
+        incarnation_id = gen_dataset_ops.multi_device_iterator_init(
+            dataset._variant_tensor,  # pylint: disable=protected-access
+            self._multi_device_iterator_resource,
+            max_buffer_size=max_buffer_size)
+
+      prototype_device_datasets = []
+      for i, device in enumerate(self._devices):
+        with ops.device(device):
+          ds = _PerDeviceGenerator(
+              i,
+              self._multi_device_iterator_resource,
+              incarnation_id,
+              source_device_tensor,
+              dataset.element_spec,
+              iterator_is_anonymous=True,
+          )
+          prototype_device_datasets.append(ds)
+
+      # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to
+      # initialize the device side of the pipeline. This would allow the
+      # MultiDeviceIterator to choose, for example, to move some transformations
+      # into the device side from its input. It might be useful in rewriting.
+      # Create the per device iterators.
+      self._device_iterators = []
+
+      for i, device in enumerate(self._devices):
+        with ops.device(device):
+          ds = _create_device_dataset(prototype_device_datasets[i],
+                                      incarnation_id, prefetch_buffer_size,
+                                      experimental_slack)
+          iterator = iter(ds)
+          self._device_iterators.append(iterator)
+
+  def get_next(self, device=None):
+    """Returns the next element given a `device`, else returns all in a list."""
+    if device is not None:
+      index = self._devices.index(device)
+      return self._device_iterators[index].get_next()
+
+    result = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        result.append(self._device_iterators[i].get_next())
+    return result
+
+  def __iter__(self):
+    return self
+
+  def next(self):
+    return self.__next__()
+
+  def __next__(self):
+    try:
+      return self.get_next()
+    except errors.OutOfRangeError:
+      raise StopIteration
+
+  def get_next_as_optional(self):
+    result = []
+    for i, device in enumerate(self._devices):
+      with ops.device(device):
+        result.append(self._device_iterators[i].get_next_as_optional())
+    return result
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+  @property
+  def _type_spec(self):
+    return MultiDeviceIteratorSpec(self._devices, self._source_device,
+                                   self._element_spec)

tensorflow/python/data/ops/optional_ops.py

@@ -0,0 +1,271 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A type for representing values that may or may not exist."""
+import abc
+
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.data.util import structure
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import gen_optional_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("experimental.Optional", "data.experimental.Optional")
+@deprecation.deprecated_endpoints("data.experimental.Optional")
+class Optional(composite_tensor.CompositeTensor, metaclass=abc.ABCMeta):
+  """Represents a value that may or may not be present.
+
+  A `tf.experimental.Optional` can represent the result of an operation that may
+  fail as a value, rather than raising an exception and halting execution. For
+  example, `tf.data.Iterator.get_next_as_optional()` returns a
+  `tf.experimental.Optional` that either contains the next element of an
+  iterator if one exists, or an "empty" value that indicates the end of the
+  sequence has been reached.
+
+  `tf.experimental.Optional` can only be used with values that are convertible
+  to `tf.Tensor` or `tf.CompositeTensor`.
+
+  One can create a `tf.experimental.Optional` from a value using the
+  `from_value()` method:
+
+  >>> optional = tf.experimental.Optional.from_value(42)
+  >>> print(optional.has_value())
+  tf.Tensor(True, shape=(), dtype=bool)
+  >>> print(optional.get_value())
+  tf.Tensor(42, shape=(), dtype=int32)
+
+  or without a value using the `empty()` method:
+
+  >>> optional = tf.experimental.Optional.empty(
+  ...   tf.TensorSpec(shape=(), dtype=tf.int32, name=None))
+  >>> print(optional.has_value())
+  tf.Tensor(False, shape=(), dtype=bool)
+  """
+
+  @abc.abstractmethod
+  def has_value(self, name=None):
+    """Returns a tensor that evaluates to `True` if this optional has a value.
+
+    >>> optional = tf.experimental.Optional.from_value(42)
+    >>> print(optional.has_value())
+    tf.Tensor(True, shape=(), dtype=bool)
+
+    Args:
+      name: (Optional.) A name for the created operation.
+
+    Returns:
+      A scalar `tf.Tensor` of type `tf.bool`.
+    """
+    raise NotImplementedError("Optional.has_value()")
+
+  @abc.abstractmethod
+  def get_value(self, name=None):
+    """Returns the value wrapped by this optional.
+
+    If this optional does not have a value (i.e. `self.has_value()` evaluates to
+    `False`), this operation will raise `tf.errors.InvalidArgumentError` at
+    runtime.
+
+    >>> optional = tf.experimental.Optional.from_value(42)
+    >>> print(optional.get_value())
+    tf.Tensor(42, shape=(), dtype=int32)
+
+    Args:
+      name: (Optional.) A name for the created operation.
+
+    Returns:
+      The wrapped value.
+    """
+    raise NotImplementedError("Optional.get_value()")
+
+  @abc.abstractproperty
+  def element_spec(self):
+    """The type specification of an element of this optional.
+
+    >>> optional = tf.experimental.Optional.from_value(42)
+    >>> print(optional.element_spec)
+    tf.TensorSpec(shape=(), dtype=tf.int32, name=None)
+
+    Returns:
+      A (nested) structure of `tf.TypeSpec` objects matching the structure of an
+      element of this optional, specifying the type of individual components.
+    """
+    raise NotImplementedError("Optional.element_spec")
+
+  @staticmethod
+  def empty(element_spec):
+    """Returns an `Optional` that has no value.
+
+    NOTE: This method takes an argument that defines the structure of the value
+    that would be contained in the returned `Optional` if it had a value.
+
+    >>> optional = tf.experimental.Optional.empty(
+    ...   tf.TensorSpec(shape=(), dtype=tf.int32, name=None))
+    >>> print(optional.has_value())
+    tf.Tensor(False, shape=(), dtype=bool)
+
+    Args:
+      element_spec: A (nested) structure of `tf.TypeSpec` objects matching the
+        structure of an element of this optional.
+
+    Returns:
+      A `tf.experimental.Optional` with no value.
+    """
+    return _OptionalImpl(gen_optional_ops.optional_none(), element_spec)
+
+  @staticmethod
+  def from_value(value):
+    """Returns a `tf.experimental.Optional` that wraps the given value.
+
+    >>> optional = tf.experimental.Optional.from_value(42)
+    >>> print(optional.has_value())
+    tf.Tensor(True, shape=(), dtype=bool)
+    >>> print(optional.get_value())
+    tf.Tensor(42, shape=(), dtype=int32)
+
+    Args:
+      value: A value to wrap. The value must be convertible to `tf.Tensor` or
+        `tf.CompositeTensor`.
+
+    Returns:
+      A `tf.experimental.Optional` that wraps `value`.
+    """
+    with ops.name_scope("optional") as scope:
+      with ops.name_scope("value"):
+        element_spec = structure.type_spec_from_value(value)
+        encoded_value = structure.to_tensor_list(element_spec, value)
+
+    return _OptionalImpl(
+        gen_optional_ops.optional_from_value(encoded_value, name=scope),
+        element_spec,
+    )
+
+
+class _OptionalImpl(Optional):
+  """Concrete implementation of `tf.experimental.Optional`.
+
+  NOTE(mrry): This implementation is kept private, to avoid defining
+  `Optional.__init__()` in the public API.
+  """
+
+  def __init__(self, variant_tensor, element_spec):
+    super().__init__()
+    self._variant_tensor = variant_tensor
+    self._element_spec = element_spec
+
+  def has_value(self, name=None):
+    with ops.colocate_with(self._variant_tensor):
+      return gen_optional_ops.optional_has_value(
+          self._variant_tensor, name=name
+      )
+
+  def get_value(self, name=None):
+    # TODO(b/110122868): Consolidate the restructuring logic with similar logic
+    # in `Iterator.get_next()` and `StructuredFunctionWrapper`.
+    with ops.name_scope(name, "OptionalGetValue",
+                        [self._variant_tensor]) as scope:
+      with ops.colocate_with(self._variant_tensor):
+        result = gen_optional_ops.optional_get_value(
+            self._variant_tensor,
+            name=scope,
+            output_types=structure.get_flat_tensor_types(self._element_spec),
+            output_shapes=structure.get_flat_tensor_shapes(self._element_spec),
+        )
+      # NOTE: We do not colocate the deserialization of composite tensors
+      # because not all ops are guaranteed to have non-GPU kernels.
+      return structure.from_tensor_list(self._element_spec, result)
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+  @property
+  def _type_spec(self):
+    return OptionalSpec.from_value(self)
+
+
+@tf_export(
+    "OptionalSpec", v1=["OptionalSpec", "data.experimental.OptionalStructure"])
+class OptionalSpec(type_spec.TypeSpec):
+  """Type specification for `tf.experimental.Optional`.
+
+  For instance, `tf.OptionalSpec` can be used to define a tf.function that takes
+  `tf.experimental.Optional` as an input argument:
+
+  >>> @tf.function(input_signature=[tf.OptionalSpec(
+  ...   tf.TensorSpec(shape=(), dtype=tf.int32, name=None))])
+  ... def maybe_square(optional):
+  ...   if optional.has_value():
+  ...     x = optional.get_value()
+  ...     return x * x
+  ...   return -1
+  >>> optional = tf.experimental.Optional.from_value(5)
+  >>> print(maybe_square(optional))
+  tf.Tensor(25, shape=(), dtype=int32)
+
+  Attributes:
+    element_spec: A (nested) structure of `TypeSpec` objects that represents the
+      type specification of the optional element.
+  """
+
+  __slots__ = ["_element_spec"]
+
+  def __init__(self, element_spec):
+    super().__init__()
+    self._element_spec = element_spec
+
+  @property
+  def value_type(self):
+    return _OptionalImpl
+
+  def _serialize(self):
+    return (self._element_spec,)
+
+  @property
+  def _component_specs(self):
+    return [tensor_spec.TensorSpec((), dtypes.variant)]
+
+  def _to_components(self, value):
+    return [value._variant_tensor]  # pylint: disable=protected-access
+
+  def _from_components(self, flat_value):
+    # pylint: disable=protected-access
+    return _OptionalImpl(flat_value[0], self._element_spec)
+
+  @staticmethod
+  def from_value(value):
+    return OptionalSpec(value.element_spec)
+
+  def _to_legacy_output_types(self):
+    return self
+
+  def _to_legacy_output_shapes(self):
+    return self
+
+  def _to_legacy_output_classes(self):
+    return self
+
+
+nested_structure_coder.register_codec(
+    nested_structure_coder.BuiltInTypeSpecCodec(
+        OptionalSpec, struct_pb2.TypeSpecProto.OPTIONAL_SPEC
+    )
+)

tensorflow/python/data/ops/options.py

@@ -0,0 +1,825 @@
+# 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.
+# ==============================================================================
+"""API for specifying `tf.data` options."""
+
+import enum
+import platform
+
+from absl import logging
+
+from tensorflow.core.framework import dataset_options_pb2
+from tensorflow.core.framework import model_pb2
+from tensorflow.python.data.ops import test_mode
+from tensorflow.python.data.util import options as options_lib
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("data.experimental.AutotuneAlgorithm")
+class AutotuneAlgorithm(enum.Enum):
+  """Represents the type of autotuning algorithm to use.
+
+  DEFAULT: The default behavior is implementation specific and may change over
+  time.
+
+  HILL_CLIMB: In each optimization step, this algorithm chooses the optimal
+  parameter and increases its value by 1.
+
+  GRADIENT_DESCENT: In each optimization step, this algorithm updates the
+  parameter values in the optimal direction.
+
+  MAX_PARALLELISM: Similar to HILL_CLIMB but uses a relaxed stopping condition,
+  allowing the optimization to oversubscribe the CPU.
+
+  STAGE_BASED: In each optimization step, this algorithm chooses the worst
+  bottleneck parameter and increases its value by 1.
+  """
+  DEFAULT = 0
+  HILL_CLIMB = 1
+  GRADIENT_DESCENT = 2
+  MAX_PARALLELISM = 3
+  STAGE_BASED = 4
+
+  @classmethod
+  def _to_proto(cls, obj):
+    if obj == cls.DEFAULT:
+      return model_pb2.AutotuneAlgorithm.DEFAULT
+    if obj == cls.HILL_CLIMB:
+      return model_pb2.AutotuneAlgorithm.HILL_CLIMB
+    if obj == cls.GRADIENT_DESCENT:
+      return model_pb2.AutotuneAlgorithm.GRADIENT_DESCENT
+    if obj == cls.MAX_PARALLELISM:
+      return model_pb2.AutotuneAlgorithm.MAX_PARALLELISM
+    if obj == cls.STAGE_BASED:
+      return model_pb2.AutotuneAlgorithm.STAGE_BASED
+    raise ValueError(
+        f"Invalid `obj.` Supported values include `DEFAULT`, `HILL_CLIMB` "
+        f"`GRADIENT_DESCENT`, and `STAGE_BASED`. Got {obj.name}.")
+
+  @classmethod
+  def _from_proto(cls, pb):
+    if pb == model_pb2.AutotuneAlgorithm.DEFAULT:
+      return cls.DEFAULT
+    if pb == model_pb2.AutotuneAlgorithm.HILL_CLIMB:
+      return cls.HILL_CLIMB
+    if pb == model_pb2.AutotuneAlgorithm.GRADIENT_DESCENT:
+      return cls.GRADIENT_DESCENT
+    if pb == model_pb2.AutotuneAlgorithm.MAX_PARALLELISM:
+      return cls.MAX_PARALLELISM
+    if pb == model_pb2.AutotuneAlgorithm.STAGE_BASED:
+      return cls.STAGE_BASED
+    raise ValueError(
+        f"Invalid `pb.` Supported values include `DEFAULT`, `HILL_CLIMB`, "
+        f"`GRADIENT_DESCENT` and `STAGE_BASED`. Got {pb}.")
+
+
+@tf_export("data.experimental.AutoShardPolicy")
+class AutoShardPolicy(enum.IntEnum):
+  """Represents the type of auto-sharding to use.
+
+  OFF: No sharding will be performed.
+
+  AUTO: Attempts FILE-based sharding, falling back to DATA-based sharding.
+
+  FILE: Shards by input files (i.e. each worker will get a set of files to
+  process). When this option is selected, make sure that there is at least as
+  many files as workers. If there are fewer input files than workers, a runtime
+  error will be raised.
+
+  DATA: Shards by elements produced by the dataset. Each worker will process the
+  whole dataset and discard the portion that is not for itself. Note that for
+  this mode to correctly partitions the dataset elements, the dataset needs to
+  produce elements in a deterministic order.
+
+  HINT: Looks for the presence of `shard(SHARD_HINT, ...)` which is treated as a
+  placeholder to replace with `shard(num_workers, worker_index)`.
+  """
+
+  # LINT.IfChange
+  OFF = -1
+  AUTO = 0
+  FILE = 1
+  DATA = 2
+  HINT = 3
+  # LINT.ThenChange(//tensorflow/python/data/experimental/ops/data_service_ops.py:tf_data_service_sharding_policy)
+
+  @classmethod
+  def _to_proto(cls, obj):
+    """Convert enum to proto."""
+    if obj == cls.OFF:
+      return dataset_options_pb2.AutoShardPolicy.OFF
+    if obj == cls.FILE:
+      return dataset_options_pb2.AutoShardPolicy.FILE
+    if obj == cls.DATA:
+      return dataset_options_pb2.AutoShardPolicy.DATA
+    if obj == cls.AUTO:
+      return dataset_options_pb2.AutoShardPolicy.AUTO
+    if obj == cls.HINT:
+      return dataset_options_pb2.AutoShardPolicy.HINT
+    raise ValueError(
+        f"Invalid `obj.` Supported values include `OFF`, `FILE`, `DATA`,"
+        f"`AUTO`, and `HINT`. Got {obj.name}."
+    )
+
+  @classmethod
+  def _from_proto(cls, pb):
+    """Convert proto to enum."""
+    if pb == dataset_options_pb2.AutoShardPolicy.OFF:
+      return cls.OFF
+    if pb == dataset_options_pb2.AutoShardPolicy.FILE:
+      return cls.FILE
+    if pb == dataset_options_pb2.AutoShardPolicy.DATA:
+      return cls.DATA
+    if pb == dataset_options_pb2.AutoShardPolicy.AUTO:
+      return cls.AUTO
+    if pb == dataset_options_pb2.AutoShardPolicy.HINT:
+      return cls.HINT
+    raise ValueError(
+        f"Invalid `pb.` Supported values include `OFF`, `FILE`, `DATA`,"
+        f"`AUTO`, and `HINT`. Got {pb}."
+    )
+
+
+@tf_export("data.experimental.ExternalStatePolicy")
+class ExternalStatePolicy(enum.Enum):
+  """Represents how to handle external state during serialization.
+
+  See the `tf.data.Options.experimental_external_state_policy` documentation
+  for more information.
+  """
+  WARN = 0
+  IGNORE = 1
+  FAIL = 2
+
+  @classmethod
+  def _to_proto(cls, obj):
+    """Convert enum to proto."""
+    if obj == cls.IGNORE:
+      return dataset_options_pb2.ExternalStatePolicy.POLICY_IGNORE
+    if obj == cls.FAIL:
+      return dataset_options_pb2.ExternalStatePolicy.POLICY_FAIL
+    if obj == cls.WARN:
+      return dataset_options_pb2.ExternalStatePolicy.POLICY_WARN
+    raise ValueError(
+        f"Invalid `obj.` Supported values include `POLICY_IGNORE`,"
+        f"`POLICY_FAIL`, `POLICY_WARN`. Got {obj.name}.")
+
+  @classmethod
+  def _from_proto(cls, pb):
+    """Convert proto to enum."""
+    if pb == dataset_options_pb2.ExternalStatePolicy.POLICY_IGNORE:
+      return cls.IGNORE
+    if pb == dataset_options_pb2.ExternalStatePolicy.POLICY_FAIL:
+      return cls.FAIL
+    if pb == dataset_options_pb2.ExternalStatePolicy.POLICY_WARN:
+      return cls.WARN
+    raise ValueError(
+        f"Invalid `pb.` Supported values include `POLICY_IGNORE`,"
+        f"`POLICY_FAIL`, `POLICY_WARN`. Got {pb}.")
+
+
+@tf_export("data.experimental.AutotuneOptions")
+class AutotuneOptions(options_lib.OptionsBase):
+  """Represents options for autotuning dataset performance.
+
+  ```python
+  options = tf.data.Options()
+  options.autotune.enabled = False
+  dataset = dataset.with_options(options)
+  ```
+  """
+
+  enabled = options_lib.create_option(
+      name="enabled",
+      ty=bool,
+      docstring="Whether to automatically tune performance knobs. If None, "
+      "defaults to True.")
+
+  cpu_budget = options_lib.create_option(
+      name="cpu_budget",
+      ty=int,
+      docstring="When autotuning is enabled (through `autotune`), determines "
+      "the CPU budget to use. Values greater than the number of schedulable "
+      "CPU cores are allowed but may result in CPU contention. If None, "
+      "defaults to the number of schedulable CPU cores.")
+
+  ram_budget = options_lib.create_option(
+      name="ram_budget",
+      ty=int,
+      docstring="When autotuning is enabled (through `autotune`), determines "
+      "the RAM budget to use. Values greater than the available RAM in bytes "
+      "may result in OOM. If None, defaults to half of the available RAM in "
+      "bytes.")
+
+  autotune_algorithm = options_lib.create_option(
+      name="autotune_algorithm",
+      ty=AutotuneAlgorithm,
+      docstring="When autotuning is enabled (through `autotune`), determines "
+      "the algorithm to use.")
+
+  initial_parallelism = options_lib.create_option(
+      name="initial_parallelism",
+      ty=int,
+      docstring=(
+          "The initial parallelism to use for parallel transformations before"
+          " autotune has a chance to run. A higher value can help with quick"
+          " startup, but may cause the ram_budget to temporarily be exceeded."
+          " Memory-sensitive datasets should consider setting this to `1` to"
+          " avoid running out of memory. Defaults to 16."
+      ),
+  )
+
+  min_parallelism = options_lib.create_option(
+      name="min_parallelism",
+      ty=int,
+      docstring=(
+          "When true, `.map(num_parallel_calls=AUTOTUNE)` and"
+          " `.batch(num_parallel_calls=AUTOTUNE)` will be at least"
+          " parallelized by `min_parallelism` threads."
+      ),
+  )
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.AutotuneOptions()
+    if self.enabled is not None:
+      pb.enabled = self.enabled
+    if self.cpu_budget is not None:
+      pb.cpu_budget = self.cpu_budget
+    if self.ram_budget is not None:
+      pb.ram_budget = self.ram_budget
+    if self.autotune_algorithm is not None:
+      pb.autotune_algorithm = AutotuneAlgorithm._to_proto(  # pylint: disable=protected-access
+          self.autotune_algorithm)
+    if self.initial_parallelism is not None:
+      pb.initial_parallelism = self.initial_parallelism
+    if self.min_parallelism is not None:
+      pb.min_parallelism = self.min_parallelism
+    return pb
+
+  def _from_proto(self, pb):
+    if pb.WhichOneof("optional_enabled") is not None:
+      self.enabled = pb.enabled
+    if pb.WhichOneof("optional_cpu_budget") is not None:
+      self.cpu_budget = pb.cpu_budget
+    if pb.WhichOneof("optional_ram_budget") is not None:
+      self.ram_budget = pb.ram_budget
+    if pb.WhichOneof("optional_autotune_algorithm") is not None:
+      self.autotune_algorithm = AutotuneAlgorithm._from_proto(  # pylint: disable=protected-access
+          pb.autotune_algorithm)
+    if pb.WhichOneof("optional_initial_parallelism") is not None:
+      self.initial_parallelism = pb.initial_parallelism
+    if pb.WhichOneof("optional_min_parallelism") is not None:
+      self.min_parallelism = pb.min_parallelism
+
+  def _set_mutable(self, mutable):
+    """Change the mutability value to `mutable` on this options and children."""
+    # pylint: disable=protected-access
+    object.__setattr__(self, "_mutable", mutable)
+
+
+@tf_export("data.experimental.DistributeOptions")
+class DistributeOptions(options_lib.OptionsBase):
+  """Represents options for distributed data processing.
+
+  You can set the distribution options of a dataset through the
+  `experimental_distribute` property of `tf.data.Options`; the property is
+  an instance of `tf.data.experimental.DistributeOptions`.
+
+  ```python
+  options = tf.data.Options()
+  options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
+  dataset = dataset.with_options(options)
+  ```
+  """
+
+  auto_shard_policy = options_lib.create_option(
+      name="auto_shard_policy",
+      ty=AutoShardPolicy,
+      docstring="The type of sharding to use. See "
+      "`tf.data.experimental.AutoShardPolicy` for additional information.",
+      default_factory=lambda: AutoShardPolicy.AUTO)
+
+  num_devices = options_lib.create_option(
+      name="num_devices",
+      ty=int,
+      docstring=
+      "The number of devices attached to this input pipeline. This will be "
+      "automatically set by `MultiDeviceIterator`.")
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.DistributeOptions()
+    pb.auto_shard_policy = AutoShardPolicy._to_proto(self.auto_shard_policy)  # pylint: disable=protected-access
+    if self.num_devices is not None:
+      pb.num_devices = self.num_devices
+    return pb
+
+  def _from_proto(self, pb):
+    self.auto_shard_policy = AutoShardPolicy._from_proto(pb.auto_shard_policy)  # pylint: disable=protected-access
+    if pb.WhichOneof("optional_num_devices") is not None:
+      self.num_devices = pb.num_devices
+
+
+@tf_export("data.experimental.OptimizationOptions")
+class OptimizationOptions(options_lib.OptionsBase):
+  """Represents options for dataset optimizations.
+
+  You can set the optimization options of a dataset through the
+  `experimental_optimization` property of `tf.data.Options`; the property is
+  an instance of `tf.data.experimental.OptimizationOptions`.
+
+  ```python
+  options = tf.data.Options()
+  options.experimental_optimization.noop_elimination = True
+  options.experimental_optimization.apply_default_optimizations = False
+  dataset = dataset.with_options(options)
+  ```
+  """
+  apply_default_optimizations = options_lib.create_option(
+      name="apply_default_optimizations",
+      ty=bool,
+      docstring=
+      "Whether to apply default graph optimizations. If False, only graph "
+      "optimizations that have been explicitly enabled will be applied.")
+
+  filter_fusion = options_lib.create_option(
+      name="filter_fusion",
+      ty=bool,
+      docstring=
+      "Whether to fuse filter transformations. If None, defaults to False.")
+
+  filter_parallelization = options_lib.create_option(
+      name="filter_parallelization",
+      ty=bool,
+      docstring=
+      "Whether to parallelize stateless filter transformations. If None, "
+      "defaults to False.")
+
+  inject_prefetch = options_lib.create_option(
+      name="inject_prefetch",
+      ty=bool,
+      docstring=
+      "Whether to inject prefetch transformation as the last transformation "
+      "when the last transformation is a synchronous transformation. If None, "
+      "defaults to True.")
+
+  seq_interleave_prefetch = options_lib.create_option(
+      name="seq_interleave_prefetch",
+      ty=bool,
+      docstring=(
+          "Whether to replace parallel interleave using a sequential interleave"
+          " that prefetches elements from its input iterators. If None,"
+          " defaults to False."
+      ),
+  )
+
+  map_and_batch_fusion = options_lib.create_option(
+      name="map_and_batch_fusion",
+      ty=bool,
+      docstring=
+      "Whether to fuse map and batch transformations. If None, defaults to "
+      "True.")
+
+  map_and_filter_fusion = options_lib.create_option(
+      name="map_and_filter_fusion",
+      ty=bool,
+      docstring=
+      "Whether to fuse map and filter transformations. If None, defaults to "
+      "False.")
+
+  map_fusion = options_lib.create_option(
+      name="map_fusion",
+      ty=bool,
+      docstring=(
+          "Whether to fuse map transformations with `num_parallel_calls` set to"
+          " `tf.data.AUTOTUNE`, no captured inputs and same `deterministic`"
+          " value. If None, defaults to False."
+      ),
+  )
+
+  map_parallelization = options_lib.create_option(
+      name="map_parallelization",
+      ty=bool,
+      docstring=
+      "Whether to parallelize stateless map transformations. If None, defaults "
+      "to True.")
+
+  noop_elimination = options_lib.create_option(
+      name="noop_elimination",
+      ty=bool,
+      docstring=
+      "Whether to eliminate no-op transformations. If None, defaults to True.")
+
+  parallel_batch = options_lib.create_option(
+      name="parallel_batch",
+      ty=bool,
+      docstring="Whether to parallelize copying of batch elements. If None, "
+      "defaults to True.")
+
+  shuffle_and_repeat_fusion = options_lib.create_option(
+      name="shuffle_and_repeat_fusion",
+      ty=bool,
+      docstring="Whether to fuse shuffle and repeat transformations. If None, "
+      "defaults to True.")
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.OptimizationOptions()
+    if self.apply_default_optimizations is not None:
+      pb.apply_default_optimizations = self.apply_default_optimizations
+    if self.filter_fusion is not None:
+      pb.filter_fusion = self.filter_fusion
+    if self.filter_parallelization is not None:
+      pb.filter_parallelization = self.filter_parallelization
+    if self.inject_prefetch is not None:
+      pb.inject_prefetch = self.inject_prefetch
+    if self.seq_interleave_prefetch is not None:
+      pb.seq_interleave_prefetch = self.seq_interleave_prefetch
+    if self.map_and_batch_fusion is not None:
+      pb.map_and_batch_fusion = self.map_and_batch_fusion
+    if self.map_and_filter_fusion is not None:
+      pb.map_and_filter_fusion = self.map_and_filter_fusion
+    if self.map_fusion is not None:
+      pb.map_fusion = self.map_fusion
+    if self.map_parallelization is not None:
+      pb.map_parallelization = self.map_parallelization
+    if self.noop_elimination is not None:
+      pb.noop_elimination = self.noop_elimination
+    if self.parallel_batch is not None:
+      pb.parallel_batch = self.parallel_batch
+    if self.shuffle_and_repeat_fusion is not None:
+      pb.shuffle_and_repeat_fusion = self.shuffle_and_repeat_fusion
+    return pb
+
+  def _from_proto(self, pb):
+    if pb.WhichOneof("optional_apply_default_optimizations") is not None:
+      self.apply_default_optimizations = pb.apply_default_optimizations
+    if pb.WhichOneof("optional_filter_fusion") is not None:
+      self.filter_fusion = pb.filter_fusion
+    if pb.WhichOneof("optional_filter_parallelization") is not None:
+      self.filter_parallelization = pb.filter_parallelization
+    if pb.WhichOneof("optional_inject_prefetch") is not None:
+      self.inject_prefetch = pb.inject_prefetch
+    if pb.WhichOneof("optional_seq_interleave_prefetch") is not None:
+      self.seq_interleave_prefetch = pb.seq_interleave_prefetch
+    if pb.WhichOneof("optional_map_and_batch_fusion") is not None:
+      self.map_and_batch_fusion = pb.map_and_batch_fusion
+    if pb.WhichOneof("optional_map_and_filter_fusion") is not None:
+      self.map_and_filter_fusion = pb.map_and_filter_fusion
+    if pb.WhichOneof("optional_map_fusion") is not None:
+      self.map_fusion = pb.map_fusion
+    if pb.WhichOneof("optional_map_parallelization") is not None:
+      self.map_parallelization = pb.map_parallelization
+    if pb.WhichOneof("optional_noop_elimination") is not None:
+      self.noop_elimination = pb.noop_elimination
+    if pb.WhichOneof("optional_parallel_batch") is not None:
+      self.parallel_batch = pb.parallel_batch
+    if pb.WhichOneof("optional_shuffle_and_repeat_fusion") is not None:
+      self.shuffle_and_repeat_fusion = pb.shuffle_and_repeat_fusion
+
+  def _set_mutable(self, mutable):
+    """Change the mutability value to `mutable` on this options and children."""
+    # pylint: disable=protected-access
+    object.__setattr__(self, "_mutable", mutable)
+
+
+@tf_export("data.experimental.ServiceOptions")
+class ServiceOptions(options_lib.OptionsBase):
+  """Represents options for tf.data service.
+
+  You can set the service options of a dataset through the
+  `experimental_service` property of `tf.data.Options`; the property is an
+  instance of `tf.data.experimental.ServiceOptions`.
+
+  ```python
+  options = tf.data.Options()
+  options.experimental_service.pinned = True
+  dataset = dataset.with_options(options)
+  ```
+  """
+
+  pinned = options_lib.create_option(
+      name="pinned",
+      ty=bool,
+      docstring=(
+          "If true, the tf.data service client allocates data to pinned memory,"
+          " which facilitates more efficient copying from host memory to GPU"
+          " memory downstream. For gRPC, compression must be disabled for this"
+          " to take effect. For alternative data transfer protocols, this may"
+          " or may not take effect, depending on the implementation."
+      ),
+  )
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.ServiceOptions()
+    if self.pinned is not None:
+      pb.pinned = self.pinned
+    return pb
+
+  def _from_proto(self, pb):
+    if pb.WhichOneof("optional_pinned") is not None:
+      self.pinned = pb.pinned
+
+
+@deprecation.deprecated_endpoints("data.experimental.ThreadingOptions")
+@tf_export("data.experimental.ThreadingOptions", "data.ThreadingOptions")
+class ThreadingOptions(options_lib.OptionsBase):
+  """Represents options for dataset threading.
+
+  You can set the threading options of a dataset through the
+  `threading` property of `tf.data.Options`; the property is
+  an instance of `tf.data.ThreadingOptions`.
+
+  ```python
+  options = tf.data.Options()
+  options.threading.private_threadpool_size = 10
+  dataset = dataset.with_options(options)
+  ```
+  """
+
+  max_intra_op_parallelism = options_lib.create_option(
+      name="max_intra_op_parallelism",
+      ty=int,
+      docstring=
+      "If set, it overrides the maximum degree of intra-op parallelism.")
+
+  private_threadpool_size = options_lib.create_option(
+      name="private_threadpool_size",
+      ty=int,
+      docstring=
+      "If set, the dataset will use a private threadpool of the given size. "
+      "The value 0 can be used to indicate that the threadpool size should be "
+      "determined at runtime based on the number of available CPU cores.")
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.ThreadingOptions()
+    if self.max_intra_op_parallelism is not None:
+      pb.max_intra_op_parallelism = self.max_intra_op_parallelism
+    if self.private_threadpool_size is not None:
+      pb.private_threadpool_size = self.private_threadpool_size
+    return pb
+
+  def _from_proto(self, pb):
+    if pb.WhichOneof("optional_max_intra_op_parallelism") is not None:
+      self.max_intra_op_parallelism = pb.max_intra_op_parallelism
+    if pb.WhichOneof("optional_private_threadpool_size") is not None:
+      self.private_threadpool_size = pb.private_threadpool_size
+
+
+@tf_export("data.Options")
+class Options(options_lib.OptionsBase):
+  """Represents options for `tf.data.Dataset`.
+
+  A `tf.data.Options` object can be, for instance, used to control which static
+  optimizations to apply to the input pipeline graph or whether to use
+  performance modeling to dynamically tune the parallelism of operations such as
+  `tf.data.Dataset.map` or `tf.data.Dataset.interleave`.
+
+  The options are set for the entire dataset and are carried over to datasets
+  created through tf.data transformations.
+
+  The options can be set by constructing an `Options` object and using the
+  `tf.data.Dataset.with_options(options)` transformation, which returns a
+  dataset with the options set.
+
+  >>> dataset = tf.data.Dataset.range(42)
+  >>> options = tf.data.Options()
+  >>> options.deterministic = False
+  >>> dataset = dataset.with_options(options)
+  >>> print(dataset.options().deterministic)
+  False
+
+  Note: A known limitation of the `tf.data.Options` implementation is that the
+  options are not preserved across tf.function boundaries. In particular, to
+  set options for a dataset that is iterated within a tf.function, the options
+  need to be set within the same tf.function.
+  """
+
+  autotune = options_lib.create_option(
+      name="autotune",
+      ty=AutotuneOptions,
+      docstring="The autotuning options associated with the dataset. See "
+      "`tf.data.experimental.AutotuneOptions` for more details.",
+      default_factory=AutotuneOptions)
+
+  deterministic = options_lib.create_option(
+      name="deterministic",
+      ty=bool,
+      docstring=
+      "Whether the outputs need to be produced in deterministic order. If None,"
+      " defaults to True.")
+
+  experimental_deterministic = options_lib.create_option(
+      name="experimental_deterministic",
+      ty=bool,
+      docstring="DEPRECATED. Use `deterministic` instead.")
+
+  experimental_distribute = options_lib.create_option(
+      name="experimental_distribute",
+      ty=DistributeOptions,
+      docstring=
+      "The distribution strategy options associated with the dataset. See "
+      "`tf.data.experimental.DistributeOptions` for more details.",
+      default_factory=DistributeOptions)
+
+  experimental_external_state_policy = options_lib.create_option(
+      name="experimental_external_state_policy",
+      ty=ExternalStatePolicy,
+      docstring="This option can be used to override the default policy for "
+      "how to handle external state when serializing a dataset or "
+      "checkpointing its iterator. There are three settings available - "
+      "IGNORE: External state is ignored without a warning; WARN: External "
+      "state is ignored and a warning is logged; FAIL: External state results "
+      "in an error.")
+
+  experimental_optimization = options_lib.create_option(
+      name="experimental_optimization",
+      ty=OptimizationOptions,
+      docstring=
+      "The optimization options associated with the dataset. See "
+      "`tf.data.experimental.OptimizationOptions` for more details.",
+      default_factory=OptimizationOptions)
+
+  experimental_slack = options_lib.create_option(
+      name="experimental_slack",
+      ty=bool,
+      docstring="Whether to introduce 'slack' in the last `prefetch` of the "
+      "input pipeline, if it exists. This may reduce CPU contention with "
+      "accelerator host-side activity at the start of a step. The slack "
+      "frequency is determined by the number of devices attached to this "
+      "input pipeline. If None, defaults to False.")
+
+  experimental_symbolic_checkpoint = options_lib.create_option(
+      name="experimental_symbolic_checkpoint",
+      ty=bool,
+      docstring="Whether to checkpoint internal input pipeline state "
+      "maintaining cursors into data sources that identify last "
+      "element(s) produced as output to the tf.data consumer. This "
+      "is alternative to the default 'explicit' checkpointing which "
+      "stores the internal input pipeline state in the checkpoint. "
+      "Note that symbolic checkpointing is not supported for "
+      "transformations that can reorder elements.")
+
+  experimental_service = options_lib.create_option(
+      name="experimental_service",
+      ty=ServiceOptions,
+      docstring=(
+          "The tf.data service options associated with the dataset. See "
+          "`tf.data.experimental.ServiceOptions` for more details."
+      ),
+      default_factory=ServiceOptions,
+  )
+
+  experimental_threading = options_lib.create_option(
+      name="experimental_threading",
+      ty=ThreadingOptions,
+      docstring="DEPRECATED. Use `threading` instead.")
+
+  experimental_warm_start = options_lib.create_option(
+      name="experimental_warm_start",
+      ty=bool,
+      docstring=(
+          "Whether to start background threads of asynchronous transformations "
+          "upon iterator creation, as opposed to during the first call to "
+          "`next()`. Defaults to `False`. "
+          "This improves the latency of the initial 'next()' calls at "
+          "the expense of requiring more memory to hold prefetched elements "
+          "between the time of iterator construction and usage."
+      ),
+      default_factory=lambda: True if test_mode.TEST_MODE else None,
+  )
+
+  dataset_name = options_lib.create_option(
+      name="dataset_name",
+      ty=str,
+      docstring="A name for the dataset, to help in debugging.")
+
+  framework_type = options_lib.create_option(
+      name="framework_type",
+      ty=list,
+      docstring="The list of frameworks that are used to generate this "
+      "pipeline, used for telemetry.")
+
+  threading = options_lib.create_option(
+      name="threading",
+      ty=ThreadingOptions,
+      docstring="The threading options associated with the dataset. See "
+      "`tf.data.ThreadingOptions` for more details.",
+      default_factory=ThreadingOptions)
+
+  def __getattribute__(self, name):
+    if name == "experimental_threading":
+      logging.warning("options.experimental_threading is deprecated. "
+                      "Use options.threading instead.")
+      return getattr(self, "threading")
+    if name == "experimental_deterministic":
+      # TODO(aaudibert): Uncomment after internal uses have been updated.
+      # logging.warning("options.experimental_deterministic is deprecated. "
+      #                 "Use options.deterministic instead.")
+      return getattr(self, "deterministic")
+    return super(Options, self).__getattribute__(name)
+
+  def __setattr__(self, name, value):
+    if name == "experimental_threading":
+      logging.warning("options.experimental_threading is deprecated. "
+                      "Use options.threading instead.")
+      super(Options, self).__setattr__("threading", value)
+      return
+    if name == "experimental_deterministic":
+      # TODO(aaudibert): Uncomment after internal uses have been updated.
+      # logging.warning("options.experimental_deterministic is deprecated. "
+      #                 "Use options.deterministic instead.")
+      super(Options, self).__setattr__("deterministic", value)
+      return
+    if name == "experimental_symbolic_checkpoint":
+      # TODO(b/276269493): Add support for MacOS.
+      if platform.system() == "Darwin":
+        logging.warning("Symbolic checkpointing is not supported on MacOS.")
+        return
+    super(Options, self).__setattr__(name, value)
+
+  def _to_proto(self):
+    pb = dataset_options_pb2.Options()
+    if self.deterministic is not None:
+      pb.deterministic = self.deterministic
+    pb.autotune_options.CopyFrom(self.autotune._to_proto())  # pylint: disable=protected-access
+    pb.distribute_options.CopyFrom(self.experimental_distribute._to_proto())  # pylint: disable=protected-access
+    if self.experimental_external_state_policy is not None:
+      pb.external_state_policy = (
+          ExternalStatePolicy._to_proto(  # pylint: disable=protected-access
+              self.experimental_external_state_policy))
+    pb.optimization_options.CopyFrom(self.experimental_optimization._to_proto())  # pylint: disable=protected-access
+    if self.experimental_slack is not None:
+      pb.slack = self.experimental_slack
+    if self.experimental_symbolic_checkpoint is not None:
+      pb.symbolic_checkpoint = self.experimental_symbolic_checkpoint
+    if self.experimental_warm_start is not None:
+      pb.warm_start = self.experimental_warm_start
+    if self.dataset_name is not None:
+      pb.dataset_name = self.dataset_name
+    if self.framework_type:
+      for framework_type in self.framework_type:
+        pb.framework_type.append(framework_type)
+    pb.service_options.CopyFrom(self.experimental_service._to_proto())  # pylint: disable=protected-access
+    pb.threading_options.CopyFrom(self.threading._to_proto())  # pylint: disable=protected-access
+    return pb
+
+  def _from_proto(self, pb):
+    if pb.WhichOneof("optional_deterministic") is not None:
+      self.deterministic = pb.deterministic
+    self.autotune._from_proto(pb.autotune_options)  # pylint: disable=protected-access
+    self.experimental_distribute._from_proto(pb.distribute_options)  # pylint: disable=protected-access
+    if pb.WhichOneof("optional_external_state_policy") is not None:
+      self.experimental_external_state_policy = (
+          ExternalStatePolicy._from_proto(  # pylint: disable=protected-access
+              pb.external_state_policy))
+    self.experimental_optimization._from_proto(pb.optimization_options)  # pylint: disable=protected-access
+    if pb.WhichOneof("optional_slack") is not None:
+      self.experimental_slack = pb.slack
+    if pb.WhichOneof("optional_symbolic_checkpoint") is not None:
+      self.experimental_symbolic_checkpoint = pb.symbolic_checkpoint
+    if pb.WhichOneof("optional_warm_start") is not None:
+      self.experimental_warm_start = pb.warm_start
+    if pb.WhichOneof("optional_dataset_name") is not None:
+      self.dataset_name = pb.dataset_name
+    if pb.framework_type:
+      self.framework_type = []
+      for framework_type in pb.framework_type:
+        self.framework_type.append(framework_type)
+    self.experimental_service._from_proto(pb.service_options)  # pylint: disable=protected-access
+    self.threading._from_proto(pb.threading_options)  # pylint: disable=protected-access
+
+  def _set_mutable(self, mutable):
+    """Change the mutability value to `mutable` on this options and children."""
+    # pylint: disable=protected-access
+    object.__setattr__(self, "_mutable", mutable)
+    self.autotune._set_mutable(mutable)
+    self.experimental_distribute._set_mutable(mutable)
+    self.experimental_optimization._set_mutable(mutable)
+    self.threading._set_mutable(mutable)
+
+  def merge(self, options):
+    """Merges itself with the given `tf.data.Options`.
+
+    If this object and the `options` to merge set an option differently, a
+    warning is generated and this object's value is updated with the `options`
+    object's value.
+
+    Args:
+      options: The `tf.data.Options` to merge with.
+
+    Returns:
+      New `tf.data.Options` object which is the result of merging self with
+      the input `tf.data.Options`.
+    """
+    return options_lib.merge_options(self, options)

tensorflow/python/data/ops/padded_batch_op.py

@@ -0,0 +1,262 @@
+# 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 implementation of `tf.data.Dataset.padded_batch`."""
+
+import numpy as np
+
+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.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import smart_cond
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _padded_batch(input_dataset,
+                  batch_size,
+                  padded_shapes=None,
+                  padding_values=None,
+                  drop_remainder=False,
+                  name=None):
+  """See `tf.data.Dataset.padded_batch` for details."""
+  if padded_shapes is None:
+    padded_shapes = dataset_ops.get_legacy_output_shapes(input_dataset)
+    for i, shape in enumerate(nest.flatten(padded_shapes)):
+      # A `tf.TensorShape` is only false if its *rank* is unknown.
+      if not shape:
+        raise ValueError(f"You must provide `padded_shapes` argument because "
+                         f"component {i} has unknown rank.")
+  return _PaddedBatchDataset(
+      input_dataset,
+      batch_size,
+      padded_shapes,
+      padding_values,
+      drop_remainder,
+      name=name)
+
+
+def _is_padded_shape_compatible_with(padded_shape, input_component_shape):
+  """Returns `True` if `input_component_shape` can be padded to `padded_shape`.
+
+  Args:
+    padded_shape: A `tf.TensorShape`.
+    input_component_shape: A `tf.TensorShape`.
+
+  Returns:
+    `True` if `input_component_shape` can be padded to `padded_shape`, otherwise
+    `False`.
+  """
+
+  if padded_shape.dims is None or input_component_shape.dims is None:
+    return True
+  if len(padded_shape.dims) != len(input_component_shape.dims):
+    return False
+  for padded_dim, input_dim in zip(padded_shape.dims,
+                                   input_component_shape.dims):
+    if (padded_dim.value is not None and input_dim.value is not None and
+        padded_dim.value < input_dim.value):
+      return False
+  return True
+
+
+def _padded_shape_to_tensor(padded_shape, input_component_shape):
+  """Converts `padded_shape` to a `tf.Tensor` representing that shape.
+
+  Args:
+    padded_shape: A shape-like object, which may be a `tf.TensorShape`, a Python
+      sequence, or a 1-D `tf.Tensor` of `tf.int64` elements.
+    input_component_shape: A `tf.TensorShape`, with which `padded_shape` must be
+      compatible.
+
+  Returns:
+    A 1-D `tf.Tensor` of `tf.int64` elements, representing `padded_shape`.
+
+  Raises:
+    ValueError: If `padded_shape` is not a shape or not compatible with
+      `input_component_shape`.
+    TypeError: If `padded_shape` is not convertible to a `tf.int64` tensor.
+  """
+  try:
+    # Try to convert the `padded_shape` to a `tf.TensorShape`
+    padded_shape_as_shape = tensor_shape.as_shape(padded_shape)
+    # We will return the "canonical" tensor representation, which uses
+    # `-1` in place of `None`.
+    ret = ops.convert_to_tensor([
+        dim if dim is not None else -1
+        for dim in padded_shape_as_shape.as_list()
+    ],
+                                dtype=dtypes.int64)
+  except (TypeError, ValueError) as e:
+    # The argument was not trivially convertible to a
+    # `tf.TensorShape`, so fall back on the conversion to tensor
+    # machinery.
+    ret = ops.convert_to_tensor(padded_shape, preferred_dtype=dtypes.int64)
+    if ret.shape.dims is not None and len(ret.shape.dims) != 1:
+      raise ValueError(
+          f"Padded shape {padded_shape} must be a `tf.int64` vector tensor, "
+          f"but its shape was {ret.shape}.") from e
+    if ret.dtype != dtypes.int64:
+      raise TypeError(
+          f"Padded shape {padded_shape} must be a `tf.int64` vector "
+          f"tensor, but its element type was {ret.dtype.name}.") from e
+    padded_shape_as_shape = tensor_util.constant_value_as_shape(ret)
+
+  if not _is_padded_shape_compatible_with(padded_shape_as_shape,
+                                          input_component_shape):
+    raise ValueError(f"The padded shape {padded_shape_as_shape} is not "
+                     f"compatible with the shape {input_component_shape} of "
+                     f"the corresponding input component.")
+
+  return ret
+
+
+def _padding_values_or_default(padding_values, input_dataset):
+  """Returns padding values with None elements replaced with default values."""
+
+  def make_zero(t):
+    if t.base_dtype == dtypes.string:
+      return ""
+    elif t.base_dtype == dtypes.variant:
+      raise TypeError("Unable to create default padding value for a component "
+                      "of type 'variant'.")
+    elif t.base_dtype == dtypes.bfloat16:
+      # Special case `bfloat16` because it is not supported by NumPy.
+      return constant_op.constant(0, dtype=dtypes.bfloat16)
+    else:
+      return np.zeros_like(t.as_numpy_dtype())
+
+  def value_or_default(value, default):
+    return default if value is None else value
+
+  default_padding = nest.map_structure(
+      make_zero, dataset_ops.get_legacy_output_types(input_dataset))
+  return nest.map_structure_up_to(padding_values, value_or_default,
+                                  padding_values, default_padding)
+
+
+def _padding_value_to_tensor(value, output_type):
+  """Converts the padding value to a tensor.
+
+  Args:
+    value: The padding value.
+    output_type: Its expected dtype.
+
+  Returns:
+    A scalar `Tensor`.
+
+  Raises:
+    ValueError: if the padding value is not a scalar.
+    TypeError: if the padding value's type does not match `output_type`.
+  """
+  value = ops.convert_to_tensor(value, name="padding_value")
+  if not value.shape.is_compatible_with(tensor_shape.TensorShape([])):
+    raise ValueError(f"Invalid `padding_values`. `padding_values` values "
+                     f"should be scalars, but got {value.shape}.")
+  if value.dtype != output_type:
+    raise TypeError(f"Invalid `padding_values`. `padding_values` values "
+                    f"type {value.dtype} does not match type {output_type} "
+                    f"of the corresponding input component.")
+  return value
+
+
+class _PaddedBatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that batches and pads contiguous elements from its input."""
+
+  def __init__(self,
+               input_dataset,
+               batch_size,
+               padded_shapes,
+               padding_values,
+               drop_remainder,
+               name=None):
+    """See `Dataset.batch()` for details."""
+    self._input_dataset = input_dataset
+
+    def check_types(component_spec):
+      if not isinstance(component_spec, tensor_spec.TensorSpec):
+        if isinstance(component_spec, dataset_ops.DatasetSpec):
+          raise TypeError(
+              "`padded_batch` is not supported for datasets of datasets")
+        raise TypeError(f"`padded_batch` is only supported for datasets that "
+                        f"produce tensor elements but type spec of elements in "
+                        f"the input dataset is not a subclass of TensorSpec: "
+                        f"`{component_spec}`.")
+
+    nest.map_structure(check_types, input_dataset.element_spec)
+    self._input_dataset = input_dataset
+    self._batch_size = ops.convert_to_tensor(
+        batch_size, dtype=dtypes.int64, name="batch_size")
+    padding_values = _padding_values_or_default(padding_values, input_dataset)
+
+    input_shapes = dataset_ops.get_legacy_output_shapes(input_dataset)
+    flat_padded_shapes = nest.flatten_up_to(input_shapes, padded_shapes)
+
+    flat_padded_shapes_as_tensors = []
+
+    for input_component_shape, padded_shape in zip(
+        nest.flatten(input_shapes), flat_padded_shapes):
+      flat_padded_shapes_as_tensors.append(
+          _padded_shape_to_tensor(padded_shape, input_component_shape))
+
+    self._padded_shapes = nest.pack_sequence_as(input_shapes,
+                                                flat_padded_shapes_as_tensors)
+
+    # If padding_values is a single element and input_shapes is a structure,
+    # "broadcast" padding_values to the same structure as input_shapes.
+    if nest.is_nested(input_shapes) and not nest.is_nested(padding_values):
+      padding_values = nest.map_structure(lambda _: padding_values,
+                                          input_shapes)
+
+    self._padding_values = nest.map_structure_up_to(
+        input_shapes, _padding_value_to_tensor, padding_values,
+        dataset_ops.get_legacy_output_types(input_dataset))
+    self._drop_remainder = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+
+    def _padded_shape_to_batch_shape(s):
+      return tensor_shape.TensorShape([
+          tensor_util.constant_value(self._batch_size)
+          if smart_cond.smart_constant_value(self._drop_remainder) else None
+      ]).concatenate(tensor_util.constant_value_as_shape(s))
+
+    output_shapes = nest.map_structure(_padded_shape_to_batch_shape,
+                                       self._padded_shapes)
+    self._structure = structure.convert_legacy_structure(
+        dataset_ops.get_legacy_output_types(self._input_dataset), output_shapes,
+        dataset_ops.get_legacy_output_classes(self._input_dataset))
+
+    self._name = name
+    # pylint: disable=protected-access
+    variant_tensor = gen_dataset_ops.padded_batch_dataset_v2(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        batch_size=self._batch_size,
+        padded_shapes=[
+            ops.convert_to_tensor(s, dtype=dtypes.int64)
+            for s in nest.flatten(self._padded_shapes)
+        ],
+        padding_values=nest.flatten(self._padding_values),
+        drop_remainder=self._drop_remainder,
+        output_shapes=structure.get_flat_tensor_shapes(self._structure),
+        metadata=self._metadata.SerializeToString())
+    super().__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/prefetch_op.py

@@ -0,0 +1,52 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.prefetch`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _prefetch(input_dataset, buffer_size, name=None):  # pylint: disable=unused-private-name
+  """See `Dataset.prefetch()` for details."""
+  if debug_mode.DEBUG_MODE:
+    return input_dataset
+  return _PrefetchDataset(input_dataset, buffer_size, name=name)
+
+
+class _PrefetchDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that asynchronously prefetches its input."""
+
+  def __init__(self, input_dataset, buffer_size, slack_period=None, name=None):
+    """See `Dataset.prefetch()` for details."""
+    self._input_dataset = input_dataset
+    if buffer_size is None:
+      buffer_size = dataset_ops.AUTOTUNE
+    self._buffer_size = ops.convert_to_tensor(
+        buffer_size, dtype=dtypes.int64, name="buffer_size")
+    self._name = name
+    # pylint: disable=protected-access
+    # We colocate the prefetch dataset with its input as this collocation only
+    # happens automatically in graph mode.
+    with ops.colocate_with(input_dataset._variant_tensor):
+      variant_tensor = gen_dataset_ops.prefetch_dataset(
+          input_dataset._variant_tensor,
+          buffer_size=self._buffer_size,
+          slack_period=slack_period,
+          legacy_autotune=(buffer_size == dataset_ops.AUTOTUNE),
+          **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/ragged_batch_op.py

@@ -0,0 +1,106 @@
+# 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 implementation of `tf.data.Dataset.ragged_batch`."""
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops.ragged import ragged_tensor
+
+
+def _ragged_batch(input_dataset,
+                  batch_size,
+                  drop_remainder=False,
+                  row_splits_dtype=dtypes.int64,
+                  name=None):
+  ragged_dataset = _DenseToRaggedDataset(input_dataset, row_splits_dtype, name)
+  return ragged_dataset.batch(batch_size, drop_remainder)
+
+
+class _DenseToRaggedDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that encodes dense inputs as ragged (w/ ragged_rank=0).
+
+  In particular:
+
+  * Any tf.Tensor elements with rank>0 are encoded as ragged tensors with
+    ragged_rank=0.  This allows tensors with varying shape to be batched
+    together.
+  * Any other elements are left as-is.
+  """
+
+  def __init__(self, input_dataset, row_splits_dtype, name=None):
+    """Constructs a new _DenseToRaggedDataset.
+
+    Args:
+      input_dataset: The dataset whose tf.Tensor elements should be made ragged.
+      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.
+      name: (Optional.) A string indicating a name for the `tf.data` operation.
+    """
+    # Replace each TensorSpec in the input dataset's structure with a
+    # corresponding RaggedTensorSpec.
+    def to_ragged_spec(spec):
+      """Returns the new spec based on RaggedTensors."""
+      if (not isinstance(spec, tensor.TensorSpec) or
+          spec.shape.rank is None or
+          spec.shape.is_fully_defined()):
+        return spec
+      else:
+        ragged_rank = max([
+            axis for (axis, size) in enumerate(spec.shape.as_list())
+            if size is None
+        ])
+        return ragged_tensor.RaggedTensorSpec(
+            shape=spec.shape,
+            dtype=spec.dtype,
+            ragged_rank=ragged_rank,
+            row_splits_dtype=row_splits_dtype)
+
+    self._structure = nest.map_structure(to_ragged_spec,
+                                         input_dataset.element_spec)
+
+    # Replace each tf.Tensor value in the input dataset with a variant-encoded
+    # RaggedTensor. Since we're updating the corresponding structure to be
+    # a RaggedTensorSpec, this variant-encoded tensor will be decoded with
+    # RaggedTensorSpec._from_tensor_list.
+    def to_ragged_variant(value):
+      """Re-encode Tensors as RaggedTensors."""
+      if (not isinstance(value, tensor.Tensor) or
+          value.shape.rank is None or
+          value.shape.is_fully_defined()):
+        return value
+      else:
+        spec = to_ragged_spec(tensor.TensorSpec.from_tensor(value))
+        if spec._ragged_rank > 0:  # pylint: disable=protected-access
+          value = ragged_tensor.RaggedTensor.from_tensor(
+              value, ragged_rank=spec._ragged_rank)  # pylint: disable=protected-access
+        return spec._to_tensor_list(value)[0]  # pylint: disable=protected-access
+
+    # Tuples are automatically unpacked by `dataset.map` so we repack them.
+    if structured_function._should_unpack(input_dataset.element_spec):  # pylint: disable=protected-access
+      map_fn = lambda *value: nest.map_structure(to_ragged_variant, value)
+    else:
+      map_fn = lambda value: nest.map_structure(to_ragged_variant, value)
+
+    self._mapped_dataset = input_dataset.map(map_fn)
+    self._name = name
+    variant = self._mapped_dataset._variant_tensor  # pylint: disable=protected-access
+    super().__init__(input_dataset, variant)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/random_op.py

@@ -0,0 +1,64 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.random`."""
+
+import warnings
+
+from tensorflow.python import tf2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _random(  # pylint: disable=unused-private-name
+    seed=None,
+    rerandomize_each_iteration=None,
+    name=None):
+  """See `Dataset.random()` for details."""
+  return _RandomDataset(
+      seed=seed,
+      rerandomize_each_iteration=rerandomize_each_iteration,
+      name=name)
+
+
+class _RandomDataset(dataset_ops.DatasetSource):
+  """A `Dataset` of pseudorandom values."""
+
+  def __init__(self, seed=None, rerandomize_each_iteration=None, name=None):
+    """A `Dataset` of pseudorandom values."""
+    self._seed, self._seed2 = random_seed.get_seed(seed)
+    self._rerandomize = rerandomize_each_iteration
+    self._name = name
+    if rerandomize_each_iteration:
+      if not tf2.enabled():
+        warnings.warn("In TF 1, the `rerandomize_each_iteration=True` option "
+                      "is only supported for repeat-based epochs.")
+      variant_tensor = ged_ops.random_dataset_v2(
+          seed=self._seed,
+          seed2=self._seed2,
+          seed_generator=gen_dataset_ops.dummy_seed_generator(),
+          rerandomize_each_iteration=self._rerandomize,
+          **self._common_args)
+    else:
+      variant_tensor = ged_ops.random_dataset(
+          seed=self._seed, seed2=self._seed2, **self._common_args)
+    super().__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.int64)

tensorflow/python/data/ops/range_op.py

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.range`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _range(*args, **kwargs):  # pylint: disable=unused-private-name
+  return _RangeDataset(*args, **kwargs)
+
+
+class _RangeDataset(dataset_ops.DatasetSource):
+  """A `Dataset` of a step separated range of values."""
+
+  def __init__(self, *args, **kwargs):
+    """See `Dataset.range()` for details."""
+    self._parse_args(*args, **kwargs)
+    self._structure = tensor_spec.TensorSpec([], self._output_type)
+    variant_tensor = gen_dataset_ops.range_dataset(
+        start=self._start,
+        stop=self._stop,
+        step=self._step,
+        **self._common_args)
+    super().__init__(variant_tensor)
+
+  def _parse_args(self, *args, **kwargs):
+    """Parses arguments according to the same rules as the `range()` builtin."""
+    if len(args) == 1:
+      self._start = self._build_tensor(0, "start")
+      self._stop = self._build_tensor(args[0], "stop")
+      self._step = self._build_tensor(1, "step")
+    elif len(args) == 2:
+      self._start = self._build_tensor(args[0], "start")
+      self._stop = self._build_tensor(args[1], "stop")
+      self._step = self._build_tensor(1, "step")
+    elif len(args) == 3:
+      self._start = self._build_tensor(args[0], "start")
+      self._stop = self._build_tensor(args[1], "stop")
+      self._step = self._build_tensor(args[2], "step")
+    else:
+      raise ValueError(f"Invalid `args`. The length of `args` should be "
+                       f"between 1 and 3 but was {len(args)}.")
+    if "output_type" in kwargs:
+      self._output_type = kwargs["output_type"]
+    else:
+      self._output_type = dtypes.int64
+    self._name = kwargs["name"] if "name" in kwargs else None
+
+  def _build_tensor(self, int64_value, name):
+    return ops.convert_to_tensor(int64_value, dtype=dtypes.int64, name=name)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/readers.py

@@ -0,0 +1,726 @@
+# 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 reader Datasets."""
+import os
+
+from tensorflow.python import tf2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import from_tensor_slices_op
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import convert
+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_spec
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_dataset_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 import nest
+from tensorflow.python.util.tf_export import tf_export
+
+_DEFAULT_READER_BUFFER_SIZE_BYTES = 256 * 1024  # 256 KB
+# The default TFRecordDataset buffer size is set to -1. The actual default is
+# set in the kernel when this value is detected.
+_DEFAULT_TF_RECORD_BUFFER_SIZE_BYTES = -1
+
+
+def _normalise_fspath(path):
+  """Convert pathlib-like objects to str (__fspath__ compatibility, PEP 519)."""
+  return os.fspath(path) if isinstance(path, os.PathLike) else path
+
+
+def _create_or_validate_filenames_dataset(filenames, name=None):
+  """Creates (or validates) a dataset of filenames.
+
+  Args:
+    filenames: Either a list or dataset of filenames. If it is a list, it is
+      convert to a dataset. If it is a dataset, its type and shape is validated.
+    name: (Optional.) A name for the tf.data operation.
+
+  Returns:
+    A dataset of filenames.
+  """
+  if isinstance(filenames, data_types.DatasetV2):
+    element_type = dataset_ops.get_legacy_output_types(filenames)
+    if element_type != dtypes.string:
+      raise TypeError(
+          "The `filenames` argument must contain `tf.string` elements. Got a "
+          f"dataset of `{element_type!r}` elements.")
+    element_shape = dataset_ops.get_legacy_output_shapes(filenames)
+    if not element_shape.is_compatible_with(tensor_shape.TensorShape([])):
+      raise TypeError(
+          "The `filenames` argument must contain `tf.string` elements of shape "
+          "[] (i.e. scalars). Got a dataset of element shape "
+          f"{element_shape!r}.")
+  else:
+    filenames = nest.map_structure(_normalise_fspath, filenames)
+    filenames = ops.convert_to_tensor(filenames, dtype_hint=dtypes.string)
+    if filenames.dtype != dtypes.string:
+      raise TypeError(
+          "The `filenames` argument must contain `tf.string` elements. Got "
+          f"`{filenames.dtype!r}` elements.")
+    filenames = array_ops.reshape(filenames, [-1], name="flat_filenames")
+    filenames = from_tensor_slices_op._TensorSliceDataset(  # pylint: disable=protected-access
+        filenames,
+        is_files=True,
+        name=name)
+  return filenames
+
+
+def _create_dataset_reader(dataset_creator,
+                           filenames,
+                           num_parallel_reads=None,
+                           name=None):
+  """Creates a dataset that reads the given files using the given reader.
+
+  Args:
+    dataset_creator: A function that takes in a single file name and returns a
+      dataset.
+    filenames: A `tf.data.Dataset` containing one or more filenames.
+    num_parallel_reads: The number of parallel reads we should do.
+    name: (Optional.) A name for the tf.data operation.
+
+  Returns:
+    A `Dataset` that reads data from `filenames`.
+  """
+
+  def read_one_file(filename):
+    filename = ops.convert_to_tensor(filename, dtypes.string, name="filename")
+    return dataset_creator(filename)
+
+  if num_parallel_reads is None:
+    return filenames.flat_map(read_one_file, name=name)
+  elif num_parallel_reads == dataset_ops.AUTOTUNE:
+    return filenames.interleave(
+        read_one_file, num_parallel_calls=num_parallel_reads, name=name)
+  else:
+    return ParallelInterleaveDataset(
+        filenames,
+        read_one_file,
+        cycle_length=num_parallel_reads,
+        block_length=1,
+        sloppy=False,
+        buffer_output_elements=None,
+        prefetch_input_elements=None,
+        name=name)
+
+
+def _get_type(value):
+  """Returns the type of `value` if it is a TypeSpec."""
+
+  if isinstance(value, type_spec.TypeSpec):
+    return value.value_type()
+  else:
+    return type(value)
+
+
+class _TextLineDataset(dataset_ops.DatasetSource):
+  """A `Dataset` comprising records from one or more text files."""
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               name=None):
+    """Creates a `TextLineDataset`.
+
+    Args:
+      filenames: A `tf.string` tensor containing one or more filenames.
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      buffer_size: (Optional.) A `tf.int64` scalar denoting the number of bytes
+        to buffer. A value of 0 results in the default buffering values chosen
+        based on the compression type.
+      name: (Optional.) A name for the tf.data operation.
+    """
+    self._filenames = filenames
+    self._compression_type = convert.optional_param_to_tensor(
+        "compression_type",
+        compression_type,
+        argument_default="",
+        argument_dtype=dtypes.string)
+    self._buffer_size = convert.optional_param_to_tensor(
+        "buffer_size",
+        buffer_size,
+        argument_default=_DEFAULT_READER_BUFFER_SIZE_BYTES)
+    self._name = name
+
+    variant_tensor = gen_dataset_ops.text_line_dataset(
+        self._filenames,
+        self._compression_type,
+        self._buffer_size,
+        metadata=self._metadata.SerializeToString())
+    super(_TextLineDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+@tf_export("data.TextLineDataset", v1=[])
+class TextLineDatasetV2(dataset_ops.DatasetSource):
+  r"""Creates a `Dataset` comprising lines from one or more text files.
+
+  The `tf.data.TextLineDataset` loads text from text files and creates a dataset
+  where each line of the files becomes an element of the dataset.
+
+  For example, suppose we have 2 files "text_lines0.txt" and "text_lines1.txt"
+  with the following lines:
+
+  >>> with open('/tmp/text_lines0.txt', 'w') as f:
+  ...   f.write('the cow\n')
+  ...   f.write('jumped over\n')
+  ...   f.write('the moon\n')
+  >>> with open('/tmp/text_lines1.txt', 'w') as f:
+  ...   f.write('jack and jill\n')
+  ...   f.write('went up\n')
+  ...   f.write('the hill\n')
+
+  We can construct a TextLineDataset from them as follows:
+
+  >>> dataset = tf.data.TextLineDataset(['/tmp/text_lines0.txt',
+  ...                                    '/tmp/text_lines1.txt'])
+
+  The elements of the dataset are expected to be:
+
+  >>> for element in dataset.as_numpy_iterator():
+  ...   print(element)
+  b'the cow'
+  b'jumped over'
+  b'the moon'
+  b'jack and jill'
+  b'went up'
+  b'the hill'
+  """
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               num_parallel_reads=None,
+               name=None):
+    r"""Creates a `TextLineDataset`.
+
+    The elements of the dataset will be the lines of the input files, using
+    the newline character '\n' to denote line splits. The newline characters
+    will be stripped off of each element.
+
+    Args:
+      filenames: A `tf.data.Dataset` whose elements are `tf.string` scalars, a
+        `tf.string` tensor, or a value that can be converted to a `tf.string`
+        tensor (such as a list of Python strings).
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      buffer_size: (Optional.) A `tf.int64` scalar denoting the number of bytes
+        to buffer. A value of 0 results in the default buffering values chosen
+        based on the compression type.
+      num_parallel_reads: (Optional.) A `tf.int64` scalar representing the
+        number of files to read in parallel. If greater than one, the records of
+        files read in parallel are outputted in an interleaved order. If your
+        input pipeline is I/O bottlenecked, consider setting this parameter to a
+        value greater than one to parallelize the I/O. If `None`, files will be
+        read sequentially.
+      name: (Optional.) A name for the tf.data operation.
+    """
+    filenames = _create_or_validate_filenames_dataset(filenames, name=name)
+    self._filenames = filenames
+    self._compression_type = compression_type
+    self._buffer_size = buffer_size
+
+    def creator_fn(filename):
+      return _TextLineDataset(
+          filename, compression_type, buffer_size, name=name)
+
+    self._impl = _create_dataset_reader(
+        creator_fn, filenames, num_parallel_reads, name=name)
+    variant_tensor = self._impl._variant_tensor  # pylint: disable=protected-access
+
+    super(TextLineDatasetV2, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+@tf_export(v1=["data.TextLineDataset"])
+class TextLineDatasetV1(dataset_ops.DatasetV1Adapter):
+  """A `Dataset` comprising lines from one or more text files."""
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               num_parallel_reads=None,
+               name=None):
+    wrapped = TextLineDatasetV2(filenames, compression_type, buffer_size,
+                                num_parallel_reads, name)
+    super(TextLineDatasetV1, self).__init__(wrapped)
+
+  __init__.__doc__ = TextLineDatasetV2.__init__.__doc__
+
+  @property
+  def _filenames(self):
+    return self._dataset._filenames  # pylint: disable=protected-access
+
+  @_filenames.setter
+  def _filenames(self, value):
+    self._dataset._filenames = value  # pylint: disable=protected-access
+
+
+class _TFRecordDataset(dataset_ops.DatasetSource):
+  """A `Dataset` comprising records from one or more TFRecord files."""
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               name=None):
+    """Creates a `TFRecordDataset`.
+
+    Args:
+      filenames: A `tf.string` tensor containing one or more filenames.
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      buffer_size: (Optional.) A `tf.int64` scalar representing the number of
+        bytes in the read buffer. 0 means no buffering.
+      name: (Optional.) A name for the tf.data operation.
+    """
+    self._filenames = filenames
+    self._compression_type = convert.optional_param_to_tensor(
+        "compression_type",
+        compression_type,
+        argument_default="",
+        argument_dtype=dtypes.string)
+    self._buffer_size = convert.optional_param_to_tensor(
+        "buffer_size",
+        buffer_size,
+        argument_default=_DEFAULT_TF_RECORD_BUFFER_SIZE_BYTES)
+    self._name = name
+
+    variant_tensor = gen_dataset_ops.tf_record_dataset(
+        self._filenames, self._compression_type, self._buffer_size,
+        metadata=self._metadata.SerializeToString())
+    super(_TFRecordDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+class ParallelInterleaveDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over its input and flattens the result."""
+
+  def __init__(self,
+               input_dataset,
+               map_func,
+               cycle_length,
+               block_length,
+               sloppy,
+               buffer_output_elements,
+               prefetch_input_elements,
+               name=None):
+    """See `tf.data.experimental.parallel_interleave()` for details."""
+    self._input_dataset = input_dataset
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func, self._transformation_name(), dataset=input_dataset)
+    if not isinstance(self._map_func.output_structure, dataset_ops.DatasetSpec):
+      raise TypeError(
+          "The `map_func` argument must return a `Dataset` object. Got "
+          f"{_get_type(self._map_func.output_structure)!r}.")
+    self._element_spec = self._map_func.output_structure._element_spec  # pylint: disable=protected-access
+    self._cycle_length = ops.convert_to_tensor(
+        cycle_length, dtype=dtypes.int64, name="cycle_length")
+    self._block_length = ops.convert_to_tensor(
+        block_length, dtype=dtypes.int64, name="block_length")
+    self._buffer_output_elements = convert.optional_param_to_tensor(
+        "buffer_output_elements",
+        buffer_output_elements,
+        argument_default=2 * block_length)
+    self._prefetch_input_elements = convert.optional_param_to_tensor(
+        "prefetch_input_elements",
+        prefetch_input_elements,
+        argument_default=2 * cycle_length)
+    if sloppy is None:
+      self._deterministic = "default"
+    elif sloppy:
+      self._deterministic = "false"
+    else:
+      self._deterministic = "true"
+    self._name = name
+
+    variant_tensor = ged_ops.legacy_parallel_interleave_dataset_v2(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        self._cycle_length,
+        self._block_length,
+        self._buffer_output_elements,
+        self._prefetch_input_elements,
+        f=self._map_func.function,
+        deterministic=self._deterministic,
+        **self._common_args)
+    super(ParallelInterleaveDataset, self).__init__(input_dataset,
+                                                    variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+  def _transformation_name(self):
+    return "tf.data.experimental.parallel_interleave()"
+
+
+@tf_export("data.TFRecordDataset", v1=[])
+class TFRecordDatasetV2(dataset_ops.DatasetV2):
+  """A `Dataset` comprising records from one or more TFRecord files.
+
+  This dataset loads TFRecords from the files as bytes, exactly as they were
+  written.`TFRecordDataset` does not do any parsing or decoding on its own.
+  Parsing and decoding can be done by applying `Dataset.map` transformations
+  after the `TFRecordDataset`.
+
+  A minimal example is given below:
+
+  >>> import tempfile
+  >>> example_path = os.path.join(tempfile.gettempdir(), "example.tfrecords")
+  >>> np.random.seed(0)
+
+  >>> # Write the records to a file.
+  ... with tf.io.TFRecordWriter(example_path) as file_writer:
+  ...   for _ in range(4):
+  ...     x, y = np.random.random(), np.random.random()
+  ...
+  ...     record_bytes = tf.train.Example(features=tf.train.Features(feature={
+  ...         "x": tf.train.Feature(float_list=tf.train.FloatList(value=[x])),
+  ...         "y": tf.train.Feature(float_list=tf.train.FloatList(value=[y])),
+  ...     })).SerializeToString()
+  ...     file_writer.write(record_bytes)
+
+  >>> # Read the data back out.
+  >>> def decode_fn(record_bytes):
+  ...   return tf.io.parse_single_example(
+  ...       # Data
+  ...       record_bytes,
+  ...
+  ...       # Schema
+  ...       {"x": tf.io.FixedLenFeature([], dtype=tf.float32),
+  ...        "y": tf.io.FixedLenFeature([], dtype=tf.float32)}
+  ...   )
+
+  >>> for batch in tf.data.TFRecordDataset([example_path]).map(decode_fn):
+  ...   print("x = {x:.4f},  y = {y:.4f}".format(**batch))
+  x = 0.5488,  y = 0.7152
+  x = 0.6028,  y = 0.5449
+  x = 0.4237,  y = 0.6459
+  x = 0.4376,  y = 0.8918
+  """
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               num_parallel_reads=None,
+               name=None):
+    """Creates a `TFRecordDataset` to read one or more TFRecord files.
+
+    Each element of the dataset will contain a single TFRecord.
+
+    Args:
+      filenames: A `tf.string` tensor or `tf.data.Dataset` containing one or
+        more filenames.
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      buffer_size: (Optional.) A `tf.int64` scalar representing the number of
+        bytes in the read buffer. If your input pipeline is I/O bottlenecked,
+        consider setting this parameter to a value 1-100 MBs. If `None`, a
+        sensible default for both local and remote file systems is used.
+      num_parallel_reads: (Optional.) A `tf.int64` scalar representing the
+        number of files to read in parallel. If greater than one, the records of
+        files read in parallel are outputted in an interleaved order. If your
+        input pipeline is I/O bottlenecked, consider setting this parameter to a
+        value greater than one to parallelize the I/O. If `None`, files will be
+        read sequentially.
+      name: (Optional.) A name for the tf.data operation.
+
+    Raises:
+      TypeError: If any argument does not have the expected type.
+      ValueError: If any argument does not have the expected shape.
+    """
+    filenames = _create_or_validate_filenames_dataset(filenames, name=name)
+
+    self._filenames = filenames
+    self._compression_type = compression_type
+    self._buffer_size = buffer_size
+    self._num_parallel_reads = num_parallel_reads
+
+    def creator_fn(filename):
+      return _TFRecordDataset(
+          filename, compression_type, buffer_size, name=name)
+
+    self._impl = _create_dataset_reader(
+        creator_fn, filenames, num_parallel_reads, name=name)
+    variant_tensor = self._impl._variant_tensor  # pylint: disable=protected-access
+    super(TFRecordDatasetV2, self).__init__(variant_tensor)
+
+  def _inputs(self):
+    return self._impl._inputs()  # pylint: disable=protected-access
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+@tf_export(v1=["data.TFRecordDataset"])
+class TFRecordDatasetV1(dataset_ops.DatasetV1Adapter):
+  """A `Dataset` comprising records from one or more TFRecord files."""
+
+  def __init__(self,
+               filenames,
+               compression_type=None,
+               buffer_size=None,
+               num_parallel_reads=None,
+               name=None):
+    wrapped = TFRecordDatasetV2(
+        filenames, compression_type, buffer_size, num_parallel_reads, name=name)
+    super(TFRecordDatasetV1, self).__init__(wrapped)
+
+  __init__.__doc__ = TFRecordDatasetV2.__init__.__doc__
+
+  @property
+  def _filenames(self):
+    return self._dataset._filenames  # pylint: disable=protected-access
+
+  @_filenames.setter
+  def _filenames(self, value):
+    self._dataset._filenames = value  # pylint: disable=protected-access
+
+
+class _FixedLengthRecordDataset(dataset_ops.DatasetSource):
+  """A `Dataset` of fixed-length records from one or more binary files."""
+
+  def __init__(self,
+               filenames,
+               record_bytes,
+               header_bytes=None,
+               footer_bytes=None,
+               buffer_size=None,
+               compression_type=None,
+               name=None):
+    """Creates a `FixedLengthRecordDataset`.
+
+    Args:
+      filenames: A `tf.string` tensor containing one or more filenames.
+      record_bytes: A `tf.int64` scalar representing the number of bytes in each
+        record.
+      header_bytes: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to skip at the start of a file.
+      footer_bytes: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to ignore at the end of a file.
+      buffer_size: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to buffer when reading.
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      name: (Optional.) A name for the tf.data operation.
+    """
+    self._filenames = filenames
+    self._record_bytes = ops.convert_to_tensor(
+        record_bytes, dtype=dtypes.int64, name="record_bytes")
+    self._header_bytes = convert.optional_param_to_tensor(
+        "header_bytes", header_bytes)
+    self._footer_bytes = convert.optional_param_to_tensor(
+        "footer_bytes", footer_bytes)
+    self._buffer_size = convert.optional_param_to_tensor(
+        "buffer_size", buffer_size, _DEFAULT_READER_BUFFER_SIZE_BYTES)
+    self._compression_type = convert.optional_param_to_tensor(
+        "compression_type",
+        compression_type,
+        argument_default="",
+        argument_dtype=dtypes.string)
+    self._name = name
+
+    variant_tensor = gen_dataset_ops.fixed_length_record_dataset_v2(
+        self._filenames,
+        self._header_bytes,
+        self._record_bytes,
+        self._footer_bytes,
+        self._buffer_size,
+        self._compression_type,
+        metadata=self._metadata.SerializeToString())
+    super(_FixedLengthRecordDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+@tf_export("data.FixedLengthRecordDataset", v1=[])
+class FixedLengthRecordDatasetV2(dataset_ops.DatasetSource):
+  """A `Dataset` of fixed-length records from one or more binary files.
+
+  The `tf.data.FixedLengthRecordDataset` reads fixed length records from binary
+  files and creates a dataset where each record becomes an element of the
+  dataset. The binary files can have a fixed length header and a fixed length
+  footer, which will both be skipped.
+
+  For example, suppose we have 2 files "fixed_length0.bin" and
+  "fixed_length1.bin" with the following content:
+
+  >>> with open('/tmp/fixed_length0.bin', 'wb') as f:
+  ...   f.write(b'HEADER012345FOOTER')
+  >>> with open('/tmp/fixed_length1.bin', 'wb') as f:
+  ...   f.write(b'HEADER6789abFOOTER')
+
+  We can construct a `FixedLengthRecordDataset` from them as follows:
+
+  >>> dataset1 = tf.data.FixedLengthRecordDataset(
+  ...     filenames=['/tmp/fixed_length0.bin', '/tmp/fixed_length1.bin'],
+  ...     record_bytes=2, header_bytes=6, footer_bytes=6)
+
+  The elements of the dataset are:
+
+  >>> for element in dataset1.as_numpy_iterator():
+  ...   print(element)
+  b'01'
+  b'23'
+  b'45'
+  b'67'
+  b'89'
+  b'ab'
+  """
+
+  def __init__(self,
+               filenames,
+               record_bytes,
+               header_bytes=None,
+               footer_bytes=None,
+               buffer_size=None,
+               compression_type=None,
+               num_parallel_reads=None,
+               name=None):
+    """Creates a `FixedLengthRecordDataset`.
+
+    Args:
+      filenames: A `tf.string` tensor or `tf.data.Dataset` containing one or
+        more filenames.
+      record_bytes: A `tf.int64` scalar representing the number of bytes in each
+        record.
+      header_bytes: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to skip at the start of a file.
+      footer_bytes: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to ignore at the end of a file.
+      buffer_size: (Optional.) A `tf.int64` scalar representing the number of
+        bytes to buffer when reading.
+      compression_type: (Optional.) A `tf.string` scalar evaluating to one of
+        `""` (no compression), `"ZLIB"`, or `"GZIP"`.
+      num_parallel_reads: (Optional.) A `tf.int64` scalar representing the
+        number of files to read in parallel. If greater than one, the records of
+        files read in parallel are outputted in an interleaved order. If your
+        input pipeline is I/O bottlenecked, consider setting this parameter to a
+        value greater than one to parallelize the I/O. If `None`, files will be
+        read sequentially.
+      name: (Optional.) A name for the tf.data operation.
+    """
+    filenames = _create_or_validate_filenames_dataset(filenames, name=name)
+
+    self._filenames = filenames
+    self._record_bytes = record_bytes
+    self._header_bytes = header_bytes
+    self._footer_bytes = footer_bytes
+    self._buffer_size = buffer_size
+    self._compression_type = compression_type
+
+    def creator_fn(filename):
+      return _FixedLengthRecordDataset(
+          filename,
+          record_bytes,
+          header_bytes,
+          footer_bytes,
+          buffer_size,
+          compression_type,
+          name=name)
+
+    self._impl = _create_dataset_reader(
+        creator_fn, filenames, num_parallel_reads, name=name)
+    variant_tensor = self._impl._variant_tensor  # pylint: disable=protected-access
+    super(FixedLengthRecordDatasetV2, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return tensor_spec.TensorSpec([], dtypes.string)
+
+
+@tf_export(v1=["data.FixedLengthRecordDataset"])
+class FixedLengthRecordDatasetV1(dataset_ops.DatasetV1Adapter):
+  """A `Dataset` of fixed-length records from one or more binary files."""
+
+  def __init__(self,
+               filenames,
+               record_bytes,
+               header_bytes=None,
+               footer_bytes=None,
+               buffer_size=None,
+               compression_type=None,
+               num_parallel_reads=None,
+               name=None):
+    wrapped = FixedLengthRecordDatasetV2(
+        filenames,
+        record_bytes,
+        header_bytes,
+        footer_bytes,
+        buffer_size,
+        compression_type,
+        num_parallel_reads,
+        name=name)
+    super(FixedLengthRecordDatasetV1, self).__init__(wrapped)
+
+  __init__.__doc__ = FixedLengthRecordDatasetV2.__init__.__doc__
+
+  @property
+  def _filenames(self):
+    return self._dataset._filenames  # pylint: disable=protected-access
+
+  @_filenames.setter
+  def _filenames(self, value):
+    self._dataset._filenames = value  # pylint: disable=protected-access
+
+
+if tf2.enabled():
+  FixedLengthRecordDataset = FixedLengthRecordDatasetV2
+  TFRecordDataset = TFRecordDatasetV2
+  TextLineDataset = TextLineDatasetV2
+else:
+  FixedLengthRecordDataset = FixedLengthRecordDatasetV1
+  TFRecordDataset = TFRecordDatasetV1
+  TextLineDataset = TextLineDatasetV1
+
+
+def _tf2_callback():
+  global FixedLengthRecordDataset, TFRecordDataset, TextLineDataset
+  if tf2.enabled():
+    FixedLengthRecordDataset = FixedLengthRecordDatasetV2
+    TFRecordDataset = TFRecordDatasetV2
+    TextLineDataset = TextLineDatasetV2
+  else:
+    FixedLengthRecordDataset = FixedLengthRecordDatasetV1
+    TFRecordDataset = TFRecordDatasetV1
+    TextLineDataset = TextLineDatasetV1
+
+
+v2_compat.register_data_v2_callback(_tf2_callback)

tensorflow/python/data/ops/rebatch_op.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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.rebatch`."""
+
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+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 gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.util import numpy_compat
+
+
+def _rebatch(input_dataset, batch_size, drop_remainder=False, name=None):
+  return _RebatchDataset(input_dataset, batch_size, drop_remainder, name)
+
+
+class _RebatchDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that rebatches elements from its input into new batch sizes.
+
+  `_RebatchDataset(input_dataset, batch_sizes)` is functionally equivalent to
+  `input_dataset.unbatch().batch(N)`, where the value of N cycles through the
+  `batch_sizes` input list. The elements produced by this dataset have the same
+  rank as the elements of the input dataset.
+  """
+
+  def __init__(self,
+               input_dataset,
+               batch_sizes,
+               drop_remainder=False,
+               name=None):
+    """See `Dataset.rebatch` for details."""
+    self._input_dataset = input_dataset
+    self._batch_sizes = ops.convert_to_tensor(
+        batch_sizes, dtype=dtypes.int64, name="batch_sizes")
+    self._drop_remainder = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+    self._name = name
+    new_batch_dim = self._compute_static_batch_dim()
+
+    # pylint: disable=protected-access
+    self._element_spec = nest.map_structure(
+        lambda ts: ts._unbatch()._batch(new_batch_dim),
+        dataset_ops.get_structure(input_dataset))
+    # pylint: enable=protected-access
+
+    # 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_v2(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        batch_sizes=batch_sizes,
+        drop_remainder=drop_remainder,
+        **self._flat_structure)
+    # LINT.ThenChange(//tensorflow/core/grappler/optimizers/data/auto_shard.cc)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _compute_static_batch_dim(self):
+    """Computes the static batch dimension of a dataset if it can be determined.
+
+    Given the RebatchDataset parameters, determines the batch dimension of this
+    dataset statically. Returns None if this cannot be determined or is
+    variable.
+
+    Returns:
+      An integer representing the batch dimension of the dataset. If it cannot
+      be determined statically, returns None.
+
+    Raises:
+      ValueError: The batch_sizes parameter is malformed, input_dataset is
+      not batched, or input_dataset batch sizes are incompatible with each
+      other.
+    """
+    new_batch_dim = tensor_util.constant_value(self._batch_sizes)
+    if new_batch_dim is None:
+      return None
+
+    if isinstance(new_batch_dim, np.ndarray):
+      if len(new_batch_dim.shape) == 1:
+        if np.all(new_batch_dim == new_batch_dim[0]):
+          new_batch_dim = new_batch_dim[0]
+        else:
+          return None
+      elif len(new_batch_dim.shape) > 1:
+        raise ValueError(
+            f"Invalid `batch_sizes`. Expected `batch_sizes` to be a scalar or "
+            f"a vector. Received `batch_sizes` of rank "
+            f"{len(new_batch_dim.shape)}.")
+
+    if self._may_form_partial_batches(new_batch_dim):
+      return None
+
+    return new_batch_dim
+
+  def _may_form_partial_batches(self, desired_batch_size):
+    """Returns whether this dataset may form partial batches."""
+    if tensor_util.constant_value(self._drop_remainder):
+      return False
+
+    def get_batch_dim(type_spec):
+      try:
+        shape = type_spec._to_legacy_output_shapes()  # pylint: disable=protected-access
+      except NotImplementedError:
+        return None
+      if not isinstance(shape, tensor_shape.TensorShape):
+        return None
+      if shape.rank is None:
+        return None
+      if len(shape) < 1:
+        raise ValueError("Invalid `batch_sizes`. Expected dataset with "
+                         "rank of >= 1 but found a dataset with "
+                         "scalar elements. Fix the issue by adding the `batch` "
+                         "transformation to the dataset.")
+      return shape.dims[0].value
+
+    input_batch_dims = [
+        get_batch_dim(ts)
+        for ts in nest.flatten(dataset_ops.get_structure(self._input_dataset))
+    ]
+    known_input_batch_dims = [d for d in input_batch_dims if d is not None]
+
+    if not known_input_batch_dims:
+      return True
+
+    known_input_batch_dims = numpy_compat.np_asarray(known_input_batch_dims)
+    if not np.all(known_input_batch_dims == known_input_batch_dims[0]):
+      raise ValueError(
+          f"Invalid `input_dataset.` The batch dimension of component 0 "
+          f"is {known_input_batch_dims[0]}, while the batch dimension "
+          f"of component i is {known_input_batch_dims}.")
+
+    return known_input_batch_dims[0] % desired_batch_size != 0
+
+  @property
+  def element_spec(self):
+    return self._element_spec

tensorflow/python/data/ops/repeat_op.py

@@ -0,0 +1,44 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.repeat`."""
+
+from tensorflow.python.data.ops import dataset_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 gen_dataset_ops
+
+
+def _repeat(input_dataset, count, name):  # pylint: disable=unused-private-name
+  return _RepeatDataset(input_dataset, count, name)
+
+
+class _RepeatDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that repeats its input several times."""
+
+  def __init__(self, input_dataset, count, name=None):
+    """See `Dataset.repeat()` for details."""
+    self._input_dataset = input_dataset
+    if count is None:
+      self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count")
+    else:
+      self._count = ops.convert_to_tensor(
+          count, dtype=dtypes.int64, name="count")
+    self._name = name
+    variant_tensor = gen_dataset_ops.repeat_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        count=self._count,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/sample_from_datasets_op.py

@@ -0,0 +1,123 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.sample_from_datasets`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import directed_interleave_op
+from tensorflow.python.data.ops import map_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_stateless_random_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.types import data as data_types
+
+
+def _sample_from_datasets(datasets,  # pylint: disable=unused-private-name
+                          weights=None,
+                          seed=None,
+                          stop_on_empty_dataset=False,
+                          rerandomize_each_iteration=None):
+  """See `Dataset.sample_from_datasets()` for details."""
+
+  def _skip_datasets_with_zero_weight(datasets, weights):
+    datasets_and_weights = [(dataset, weight)
+                            for (dataset, weight) in zip(datasets, weights)
+                            if weight > 0]
+    return (zip(*datasets_and_weights) if datasets_and_weights else
+            ([datasets[0].take(0)], [1.]))
+
+  if not datasets:
+    raise ValueError("Invalid `datasets`. `datasets` should not be empty.")
+
+  if not isinstance(weights, data_types.DatasetV2):
+    if weights is None:
+      # Select inputs with uniform probability.
+      logits = [[1.0] * len(datasets)]
+
+    else:
+      if isinstance(weights, tensor.Tensor):
+        if not weights.shape.is_compatible_with([len(datasets)]):
+          raise ValueError(f"Invalid `weights`. The shape of `weights` "
+                           f"should be compatible with `[len(datasets)]` "
+                           f"but is {weights.shape}.")
+      else:
+        if len(datasets) != len(weights):
+          raise ValueError(f"Invalid `weights`. `weights` should have the "
+                           f"same length as `datasets` but got "
+                           f"`len(weights)={len(weights)}` vs. "
+                           f"`len(datasets)={len(datasets)}`.")
+
+      # Use the given `weights` as the probability of choosing the respective
+      # input.
+      if not isinstance(weights, tensor.Tensor):
+        datasets, weights = _skip_datasets_with_zero_weight(datasets, weights)
+      weights = ops.convert_to_tensor(weights, name="weights")
+      if weights.dtype not in (dtypes.float32, dtypes.float64):
+        raise TypeError(f"Invalid `weights`. `weights` type must be either "
+                        f"`tf.float32` or `tf.float64` but is "
+                        f"{weights.dtype}.")
+
+      # The `stateless_multinomial()` op expects log-probabilities, as opposed
+      # to weights.
+      logits = array_ops.expand_dims(math_ops.log(weights, name="logits"), 0)
+
+    # NOTE(mrry): We only specialize when `weights` is not a `Dataset`. When
+    # it is a `Dataset`, it is possible that evaluating it has a side effect
+    # the user depends on.
+    if len(datasets) == 1:
+      return datasets[0]
+
+    def select_dataset_constant_logits(seed):
+      return array_ops.squeeze(
+          gen_stateless_random_ops.stateless_multinomial(
+              logits, 1, seed=seed),
+          axis=[0, 1])
+
+    selector_input = map_op._MapDataset(  # pylint: disable=protected-access
+        dataset_ops.Dataset.random(
+            seed=seed,
+            rerandomize_each_iteration=rerandomize_each_iteration).batch(2),
+        select_dataset_constant_logits,
+        use_inter_op_parallelism=False)
+
+  else:  # isinstance(weights, DatasetV2)
+    # Use each element of the given `weights` dataset as the probability of
+    # choosing the respective input.
+    #
+    # The `stateless_multinomial()` op expects log-probabilities, as opposed
+    # to weights.
+    logits_ds = weights.map(lambda *p: math_ops.log(p, name="logits"))
+
+    def select_dataset_varying_logits(logits, seed):
+      return array_ops.squeeze(
+          gen_stateless_random_ops.stateless_multinomial(
+              logits, 1, seed=seed),
+          axis=[0, 1])
+
+    logits_and_seeds = dataset_ops.Dataset.zip(
+        (logits_ds,
+         dataset_ops.Dataset.random(
+             seed=seed,
+             rerandomize_each_iteration=rerandomize_each_iteration).batch(2)))
+    selector_input = map_op._MapDataset(  # pylint: disable=protected-access
+        logits_and_seeds,
+        select_dataset_varying_logits,
+        use_inter_op_parallelism=False)
+
+  return directed_interleave_op._directed_interleave(  # pylint: disable=protected-access
+      selector_input, datasets, stop_on_empty_dataset
+  )

tensorflow/python/data/ops/save_op.py

@@ -0,0 +1,118 @@
+# 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.
+# ==============================================================================
+"""Implementation of SaveDataset in Python."""
+import os
+
+from tensorflow.python.checkpoint import checkpoint as checkpoint_lib
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.platform import gfile
+# TODO(b/238903802): Use TypeSpec serialization methods directly.
+from tensorflow.python.saved_model import nested_structure_coder
+
+
+def _save(input_dataset,
+          path,
+          compression=None,
+          shard_func=None,
+          checkpoint_args=None):
+  """Implements the save function and checkpoint functionality."""
+  if context.executing_eagerly() and checkpoint_args:
+    save_dataset = _SaveDataset(input_dataset, path, shard_func, compression)
+    save_iterator = iter(save_dataset)
+
+    if "checkpoint" in checkpoint_args:
+      raise ValueError(
+          "'Invalid `checkpoint_args`. `checkpoint_args` are not allowed "
+          "to include 'checkpoint'."
+      )
+    checkpoint = checkpoint_lib.Checkpoint(iterator=save_iterator)
+    checkpoint_args["checkpoint"] = checkpoint
+    manager = checkpoint_management.CheckpointManager(**checkpoint_args)
+    checkpoint.restore(manager.latest_checkpoint)
+
+    for _ in enumerate(save_iterator):
+      if "step_counter" in checkpoint_args:
+        checkpoint_args["step_counter"].assign_add(delta=1)
+      manager.save(check_interval=True)
+  else:
+    dataset, shard_func, use_shard_func, path = set_save_dataset_attributes(
+        input_dataset, shard_func, path)
+    return ged_ops.save_dataset(
+        dataset._variant_tensor,   # pylint: disable=protected-access
+        path=path,
+        shard_func_other_args=shard_func.captured_inputs,
+        compression=compression,
+        shard_func=shard_func,
+        use_shard_func=use_shard_func)
+
+
+class _SaveDataset(dataset_ops.UnaryDataset):
+  """"A dataset that loads previously saved dataset."""
+
+  def __init__(self, dataset, path, shard_func, compression):
+    self._element_spec = dataset.element_spec
+    self._shard_func = shard_func
+    dataset, shard_func, use_shard_func, path = set_save_dataset_attributes(
+        dataset, shard_func, path)
+    variant_tensor = ged_ops.save_dataset_v2(
+        dataset._variant_tensor,  # pylint: disable=protected-access
+        path=path,
+        shard_func_other_args=shard_func.captured_inputs,
+        shard_func=shard_func,
+        use_shard_func=use_shard_func,
+        compression=compression,
+        output_types=structure.get_flat_tensor_types(dataset.element_spec),
+        output_shapes=structure.get_flat_tensor_shapes(dataset.element_spec),
+    )
+    super().__init__(dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._shard_func]
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+
+def set_save_dataset_attributes(dataset, shard_func, path):
+  """Sets parameters for SaveDatasetOp and SaveDatasetV2Op."""
+  if shard_func is None:
+    use_shard_func = False
+    shard_func = lambda *x: None  # a dummy function that will not be used
+  else:
+    use_shard_func = True
+  wrapped_func = structured_function.StructuredFunctionWrapper(
+      shard_func,
+      "save()",
+      input_structure=dataset.element_spec,
+      add_to_graph=False)
+  encoded = nested_structure_coder.encode_structure(dataset.element_spec)
+  gfile.MakeDirs(path)
+  with gfile.GFile(os.path.join(path, dataset_ops.DATASET_SPEC_FILENAME),
+                   "wb") as f:
+    f.write(encoded.SerializeToString())
+  path = ops.convert_to_tensor(path, dtype=dtypes.string, name="path")
+  shard_func = wrapped_func.function
+  shard_func.add_to_graph(ops.get_default_graph())
+  # pylint: disable=protected-access
+  dataset._apply_debug_options()
+  return dataset, shard_func, use_shard_func, path

tensorflow/python/data/ops/scan_op.py

@@ -0,0 +1,172 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.shuffle`."""
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.util.compat import collections_abc
+
+
+def scan(
+    input_dataset, initial_state, scan_func, use_default_device=None, name=None
+):
+  return _ScanDataset(
+      input_dataset, initial_state, scan_func, use_default_device, name=name)
+
+
+class _ScanDataset(dataset_ops.UnaryDataset):
+  """A dataset that scans a function across its input."""
+
+  def __init__(self,
+               input_dataset,
+               initial_state,
+               scan_func,
+               use_default_device=None,
+               name=None):
+    """See `scan()` for details."""
+    self._input_dataset = input_dataset
+    self._initial_state = structure.normalize_element(initial_state)
+
+    # Compute initial values for the state classes, shapes and types based on
+    # the initial state. The shapes may be refined by running `tf_scan_func` one
+    # or more times below.
+    self._state_structure = structure.type_spec_from_value(self._initial_state)
+
+    # Iteratively rerun the scan function until reaching a fixed point on
+    # `self._state_shapes`.
+    wrapped_func = None
+    need_to_rerun = True
+    while need_to_rerun:
+
+      wrapped_func = structured_function.StructuredFunctionWrapper(
+          scan_func,
+          self._transformation_name(),
+          input_structure=(self._state_structure, input_dataset.element_spec),
+          add_to_graph=False)
+      if not (isinstance(wrapped_func.output_types, collections_abc.Sequence)
+              and len(wrapped_func.output_types) == 2):
+        raise TypeError(f"Invalid `scan_func`. `scan_func` should return a "
+                        f"pair consisting of new state and the output value "
+                        f"but its return type is "
+                        f"{wrapped_func.output_structure}.")
+
+      # Extract and validate class information from the returned values.
+      new_state_classes, output_classes = wrapped_func.output_classes
+      self._output_classes = output_classes
+      old_state_classes = nest.map_structure(
+          lambda component_spec: component_spec._to_legacy_output_classes(),  # pylint: disable=protected-access
+          self._state_structure)
+
+      flat_new_state_classes = nest.flatten(new_state_classes)
+      flat_old_state_classes = nest.flatten(old_state_classes)
+      if len(flat_new_state_classes) != len(flat_old_state_classes):
+        raise TypeError(
+            "Invalid `scan_func`. The state returned by "
+            "`scan_func` must have the same number of elements "
+            "as the initial state. Expected "
+            f"{len(flat_old_state_classes)}, got "
+            f"{len(flat_new_state_classes)}."
+        )
+
+      for new_state_class, old_state_class in zip(
+          flat_new_state_classes, flat_old_state_classes
+      ):
+        if not issubclass(new_state_class, old_state_class):
+          raise TypeError(f"Invalid `scan_func`. The element classes for the "
+                          f"new state must match the initial state. Expected "
+                          f"{old_state_classes}, got {new_state_classes}.")
+
+      # Extract and validate type information from the returned values.
+      new_state_types, output_types = wrapped_func.output_types
+      old_state_types = nest.map_structure(
+          lambda component_spec: component_spec._to_legacy_output_types(),  # pylint: disable=protected-access
+          self._state_structure)
+      for new_state_type, old_state_type in zip(
+          nest.flatten(new_state_types), nest.flatten(old_state_types)):
+        if new_state_type != old_state_type:
+          raise TypeError(f"Invalid `scan_func`. The element types for the "
+                          f"new state must match the initial state. Expected "
+                          f"{old_state_types}, got {new_state_types}.")
+
+      # Extract shape information from the returned values.
+      new_state_shapes, output_shapes = wrapped_func.output_shapes
+      old_state_shapes = nest.map_structure(
+          lambda component_spec: component_spec._to_legacy_output_shapes(),  # pylint: disable=protected-access
+          self._state_structure)
+      self._element_spec = structure.convert_legacy_structure(
+          output_types, output_shapes, output_classes)
+
+      flat_state_shapes = nest.flatten(old_state_shapes)
+      flat_new_state_shapes = nest.flatten(new_state_shapes)
+      weakened_state_shapes = [
+          original.most_specific_compatible_shape(new)
+          for original, new in zip(flat_state_shapes, flat_new_state_shapes)
+      ]
+
+      need_to_rerun = False
+      for original_shape, weakened_shape in zip(flat_state_shapes,
+                                                weakened_state_shapes):
+        if original_shape.ndims is not None and (
+            weakened_shape.ndims is None or
+            original_shape.as_list() != weakened_shape.as_list()):
+          need_to_rerun = True
+          break
+
+      if need_to_rerun:
+        # TODO(b/110122868): Support a "most specific compatible structure"
+        # method for combining structures, to avoid using legacy structures
+        # in this method.
+        self._state_structure = structure.convert_legacy_structure(
+            old_state_types,
+            nest.pack_sequence_as(old_state_shapes, weakened_state_shapes),
+            old_state_classes)
+
+    assert wrapped_func is not None
+    self._scan_func = wrapped_func
+    self._scan_func.function.add_to_graph(ops.get_default_graph())
+
+    self._name = name
+    # pylint: disable=protected-access
+    if use_default_device is not None:
+      variant_tensor = ged_ops.scan_dataset(
+          self._input_dataset._variant_tensor,
+          structure.to_tensor_list(self._state_structure, self._initial_state),
+          self._scan_func.function.captured_inputs,
+          f=self._scan_func.function,
+          preserve_cardinality=True,
+          use_default_device=use_default_device,
+          **self._common_args)
+    else:
+      variant_tensor = ged_ops.scan_dataset(
+          self._input_dataset._variant_tensor,
+          structure.to_tensor_list(self._state_structure, self._initial_state),
+          self._scan_func.function.captured_inputs,
+          f=self._scan_func.function,
+          preserve_cardinality=True,
+          **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._scan_func]
+
+  @property
+  def element_spec(self):
+    return self._element_spec
+
+  def _transformation_name(self):
+    return "Dataset.scan()"

tensorflow/python/data/ops/shard_op.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.shard`."""
+
+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
+
+
+def _shard(input_dataset, num_shards, index, name):  # pylint: disable=unused-private-name
+  """See `Dataset.shard()` for details."""
+  return _ShardDataset(input_dataset, num_shards, index, name)
+
+
+class _ShardDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` for sharding its input."""
+
+  def __init__(self, input_dataset, num_shards, index, name):
+    """See `Dataset.shard()` for details."""
+    self._input_dataset = input_dataset
+    self._num_shards = ops.convert_to_tensor(
+        num_shards, dtype=dtypes.int64, name="num_shards")
+    self._index = ops.convert_to_tensor(index, dtype=dtypes.int64, name="index")
+    self._name = name
+    variant_tensor = gen_dataset_ops.shard_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        num_shards=self._num_shards,
+        index=self._index,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/shuffle_op.py

@@ -0,0 +1,73 @@
+# 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 implementation of `tf.data.Dataset.shuffle`."""
+from tensorflow.python import tf2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import random_seed
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _shuffle(  # pylint: disable=unused-private-name
+    input_dataset,
+    buffer_size,
+    seed=None,
+    reshuffle_each_iteration=True,
+    name=None,
+):
+  return _ShuffleDataset(
+      input_dataset, buffer_size, seed, reshuffle_each_iteration, name=name)
+
+
+class _ShuffleDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` that randomly shuffles the elements of its input."""
+
+  def __init__(
+      self,
+      input_dataset,
+      buffer_size,
+      seed=None,
+      reshuffle_each_iteration=True,
+      name=None,
+  ):
+    """See `Dataset.shuffle()` for details."""
+    self._input_dataset = input_dataset
+    self._buffer_size = ops.convert_to_tensor(
+        buffer_size, dtype=dtypes.int64, name="buffer_size")
+    self._seed, self._seed2 = random_seed.get_seed(seed)
+    self._reshuffle_each_iteration = reshuffle_each_iteration
+    self._name = name
+
+    if (tf2.enabled() and
+        (context.executing_eagerly() or ops.inside_function())):
+      variant_tensor = gen_dataset_ops.shuffle_dataset_v3(
+          input_dataset._variant_tensor,  # pylint: disable=protected-access
+          buffer_size=self._buffer_size,
+          seed=self._seed,
+          seed2=self._seed2,
+          seed_generator=gen_dataset_ops.dummy_seed_generator(),
+          reshuffle_each_iteration=self._reshuffle_each_iteration,
+          **self._common_args)
+    else:
+      variant_tensor = gen_dataset_ops.shuffle_dataset(
+          input_dataset._variant_tensor,  # pylint: disable=protected-access
+          buffer_size=self._buffer_size,
+          seed=self._seed,
+          seed2=self._seed2,
+          reshuffle_each_iteration=self._reshuffle_each_iteration,
+          **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/skip_op.py

@@ -0,0 +1,38 @@
+# 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 implementation of `tf.data.Dataset.skip`."""
+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
+
+
+def _skip(self, count, name=None):  # pylint: disable=unused-private-name
+  return _SkipDataset(self, count, name)
+
+
+class _SkipDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` skipping the first `count` elements from its input."""
+
+  def __init__(self, input_dataset, count, name=None):
+    """See `Dataset.skip()` for details."""
+    self._input_dataset = input_dataset
+    self._count = ops.convert_to_tensor(count, dtype=dtypes.int64, name="count")
+    self._name = name
+    variant_tensor = gen_dataset_ops.skip_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        count=self._count,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/snapshot_op.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.snapshot`."""
+
+import multiprocessing
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _snapshot(input_dataset,  # pylint: disable=unused-private-name
+              path,
+              compression="AUTO",
+              reader_func=None,
+              shard_func=None,
+              name=None):
+  """See `Dataset.snapshot()` for details."""
+
+  project_func = None
+  if shard_func is None:
+    input_dataset = input_dataset.enumerate(name=name)
+    # This sets the amount of parallelism based on the number of CPU cores on
+    # the machine where this Python code is executed, which may differ from
+    # the number of CPU cores where the input pipeline graph is actually
+    # executed (e.g. remote Cloud TPU workers).
+    local_shard_func = lambda index, _: index % multiprocessing.cpu_count()
+    project_func = lambda _, elem: elem
+  else:
+    local_shard_func = shard_func
+  dataset = _SnapshotDataset(
+      input_dataset=input_dataset,
+      path=path,
+      compression=compression,
+      reader_func=reader_func,
+      # This will not do the right thing where the graph is built on a
+      # different machine than the executor (e.g. Cloud TPUs).
+      shard_func=local_shard_func,
+      name=name)
+  if project_func is not None:
+    dataset = dataset.map(project_func, name=name)
+  return dataset
+
+
+class _SnapshotDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A dataset that allows saving and re-use of already processed data."""
+
+  def __init__(self,
+               input_dataset,
+               path,
+               shard_func,
+               compression=None,
+               reader_func=None,
+               pending_snapshot_expiry_seconds=None,
+               use_legacy_function=False,
+               name=None):
+
+    if reader_func is None:
+      reader_func = lambda datasets: datasets.interleave(  # pylint:disable=g-long-lambda
+          lambda x: x,
+          cycle_length=multiprocessing.cpu_count(),
+          num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    self._input_dataset = input_dataset
+    self._path = path
+    self._compression = compression
+
+    self._reader_func = structured_function.StructuredFunctionWrapper(
+        reader_func,
+        self._transformation_name() + ".reader_func",
+        # Dataset of datasets of input elements
+        input_structure=dataset_ops.DatasetSpec(
+            dataset_ops.DatasetSpec(input_dataset.element_spec)),
+        use_legacy_function=use_legacy_function)
+    self._shard_func = structured_function.StructuredFunctionWrapper(
+        shard_func,
+        self._transformation_name() + ".shard_func",
+        dataset=input_dataset,
+        use_legacy_function=use_legacy_function)
+
+    if ((not self._shard_func.output_structure.is_compatible_with(
+        tensor_spec.TensorSpec([], dtypes.int32))) and
+        (not self._shard_func.output_structure.is_compatible_with(
+            tensor_spec.TensorSpec([], dtypes.int64)))):
+      raise TypeError(f"Invalid `shard_func`. `shard_func` must return "
+                      f"`tf.int64` scalar tensor but its return type is "
+                      f"{self._shard_func.output_structure}.")
+
+    self._name = name
+    variant_tensor = ged_ops.snapshot_dataset_v2(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        path,
+        self._reader_func.function.captured_inputs,
+        self._shard_func.function.captured_inputs,
+        compression=compression,
+        reader_func=self._reader_func.function,
+        shard_func=self._shard_func.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._reader_func, self._shard_func]
+
+  def _transformation_name(self):
+    return "Dataset.snapshot()"

tensorflow/python/data/ops/sparse_batch_op.py

@@ -0,0 +1,56 @@
+# 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 implementation of `tf.data.Dataset.sparse_batch`."""
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import convert
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _sparse_batch(input_dataset, batch_size, row_shape, name=None):
+  return _DenseToSparseBatchDataset(input_dataset, batch_size, row_shape, name)
+
+
+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, name=None):
+    """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))
+    self._name = name
+    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

tensorflow/python/data/ops/structured_function.py

@@ -0,0 +1,309 @@
+# 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.
+# ==============================================================================
+"""Utilities for managing tf.data user-defined functions."""
+
+import warnings
+
+from tensorflow.python.autograph.core import ag_ctx as autograph_ctx
+from tensorflow.python.autograph.impl import api as autograph
+from tensorflow.python.data.ops import debug_mode
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import structure
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.util import function_utils
+from tensorflow.python.util import variable_utils
+
+
+def _should_pack(arg):
+  """Determines whether the caller needs to pack the argument in a tuple.
+
+  If user-defined function returns a list of tensors, `nest.flatten()` and
+  `ops.convert_to_tensor()` and would conspire to attempt to stack those tensors
+  into a single tensor because the tf.data version of `nest.flatten()` does
+  not recurse into lists. Since it is more likely that the list arose from
+  returning the result of an operation (such as `tf.numpy_function()`) that
+  returns a list of not-necessarily-stackable tensors, we treat the returned
+  value as a `tuple` instead. A user wishing to pack the return value into a
+  single tensor can use an explicit `tf.stack()` before returning.
+
+  Args:
+    arg: argument to check
+
+  Returns:
+    Indication of whether the caller needs to pack the argument in a tuple.
+  """
+  return isinstance(arg, list)
+
+
+def _should_unpack(arg):
+  """Determines whether the caller needs to unpack the argument from a tuple.
+
+  Args:
+    arg: argument to check
+
+  Returns:
+    Indication of whether the caller needs to unpack the argument from a tuple.
+  """
+  return type(arg) is tuple  # pylint: disable=unidiomatic-typecheck
+
+
+class StructuredFunctionWrapper():
+  """A function wrapper that supports structured arguments and return values."""
+
+  def __init__(self,
+               func,
+               transformation_name,
+               dataset=None,
+               input_classes=None,
+               input_shapes=None,
+               input_types=None,
+               input_structure=None,
+               add_to_graph=True,
+               use_legacy_function=False,
+               defun_kwargs=None):
+    """Creates a new `StructuredFunctionWrapper` for the given function.
+
+    Args:
+      func: A function from a (nested) structure to another (nested) structure.
+      transformation_name: Human-readable name of the transformation in which
+        this function is being instantiated, for error messages.
+      dataset: (Optional.) A `tf.data.Dataset`. If given, the structure of this
+        dataset will be assumed as the structure for `func` arguments; otherwise
+        `input_classes`, `input_shapes`, and `input_types` must be defined.
+      input_classes: (Optional.) A (nested) structure of `type`. If given, this
+        argument defines the Python types for `func` arguments.
+      input_shapes: (Optional.) A (nested) structure of `tf.TensorShape`. If
+        given, this argument defines the shapes and structure for `func`
+        arguments.
+      input_types: (Optional.) A (nested) structure of `tf.DType`. If given,
+        this argument defines the element types and structure for `func`
+        arguments.
+      input_structure: (Optional.) A `Structure` object. If given, this argument
+        defines the element types and structure for `func` arguments.
+      add_to_graph: (Optional.) If `True`, the function will be added to the
+        default graph, if it exists.
+      use_legacy_function: (Optional.) A boolean that determines whether the
+        function be created using `tensorflow.python.eager.function.defun`
+        (default behavior) or `tensorflow.python.framework.function.Defun`
+        (legacy behavior).
+      defun_kwargs: (Optional.) A dictionary mapping string argument names to
+        values. If supplied, will be passed to `function` as keyword arguments.
+
+    Raises:
+      ValueError: If an invalid combination of `dataset`, `input_classes`,
+        `input_shapes`, and `input_types` is passed.
+    """
+    # pylint: disable=protected-access
+    if input_structure is None:
+      if dataset is None:
+        if input_classes is None or input_shapes is None or input_types is None:
+          raise ValueError("Either `dataset`, `input_structure` or all of "
+                           "`input_classes`, `input_shapes`, and `input_types` "
+                           "must be specified.")
+        self._input_structure = structure.convert_legacy_structure(
+            input_types, input_shapes, input_classes)
+      else:
+        if not (input_classes is None and input_shapes is None and
+                input_types is None):
+          raise ValueError("Either `dataset`, `input_structure` or all of "
+                           "`input_classes`, `input_shapes`, and `input_types` "
+                           "must be specified.")
+        self._input_structure = dataset.element_spec
+    else:
+      if not (dataset is None and input_classes is None and
+              input_shapes is None and input_types is None):
+        raise ValueError("Either `dataset`, `input_structure`, or all of "
+                         "`input_classes`, `input_shapes`, and `input_types` "
+                         "must be specified.")
+      self._input_structure = input_structure
+
+    self._func = func
+
+    if defun_kwargs is None:
+      defun_kwargs = {}
+
+    readable_transformation_name = transformation_name.replace(
+        ".", "_")[:-2] if len(transformation_name) > 2 else ""
+
+    func_name = "_".join(
+        [readable_transformation_name,
+         function_utils.get_func_name(func)])
+    # Sanitize function name to remove symbols that interfere with graph
+    # construction.
+    for symbol in ["<", ">", "\\", "'", " "]:
+      func_name = func_name.replace(symbol, "")
+
+    ag_ctx = autograph_ctx.control_status_ctx()
+
+    def wrapper_helper(*args):
+      """Wrapper for passing nested structures to and from tf.data functions."""
+      nested_args = structure.from_compatible_tensor_list(
+          self._input_structure, args)
+      if not _should_unpack(nested_args):
+        nested_args = (nested_args,)
+      ret = autograph.tf_convert(self._func, ag_ctx)(*nested_args)
+      ret = variable_utils.convert_variables_to_tensors(ret)
+      if _should_pack(ret):
+        ret = tuple(ret)
+
+      try:
+        self._output_structure = structure.type_spec_from_value(ret)
+      except (ValueError, TypeError) as e:
+        raise TypeError(f"Unsupported return value from function passed to "
+                        f"{transformation_name}: {ret}.") from e
+      return ret
+
+    def trace_legacy_function(defun_kwargs):
+
+      @function.Defun(*structure.get_flat_tensor_types(self._input_structure),
+                      **defun_kwargs)
+      def wrapped_fn(*args):
+        ret = wrapper_helper(*args)
+        return structure.to_tensor_list(self._output_structure, ret)
+
+      return lambda: wrapped_fn
+
+    def trace_py_function(defun_kwargs):
+      # First we trace the function to infer the output structure.
+      def unused(*args):  # pylint: disable=missing-docstring,unused-variable
+        ret = wrapper_helper(*args)
+        ret = structure.to_tensor_list(self._output_structure, ret)
+        return [ops.convert_to_tensor(t) for t in ret]
+
+      func_name = defun_kwargs.pop("func_name", "unused")
+      tf_function = def_function.Function(
+          python_function=unused,
+          name=func_name,
+          input_signature=structure.get_flat_tensor_specs(
+              self._input_structure
+          ),
+          autograph=False,
+          experimental_attributes=defun_kwargs,
+      )
+
+      _ = tf_function.get_concrete_function()
+
+      def py_function_wrapper(*args):
+        nested_args = structure.from_compatible_tensor_list(
+            self._input_structure, args)
+        if not _should_unpack(nested_args):
+          nested_args = (nested_args,)
+        ret = self._func(*nested_args)
+        if _should_pack(ret):
+          ret = tuple(ret)
+        ret = structure.to_tensor_list(self._output_structure, ret)
+        return [ops.convert_to_tensor(t) for t in ret]
+
+      # Next we trace the function wrapped in `eager_py_func` to force eager
+      # execution.
+      @def_function.function(
+          input_signature=structure.get_flat_tensor_specs(
+              self._input_structure),
+          autograph=False,
+          experimental_attributes=defun_kwargs)
+      def wrapped_fn(*args):  # pylint: disable=missing-docstring
+        return script_ops.eager_py_func(
+            py_function_wrapper, args,
+            structure.get_flat_tensor_types(self._output_structure))
+
+      return wrapped_fn.get_concrete_function
+
+    def trace_tf_function(defun_kwargs):
+      # Note: wrapper_helper will apply autograph based on context.
+      def wrapped_fn(*args):  # pylint: disable=missing-docstring
+        ret = wrapper_helper(*args)
+        ret = structure.to_tensor_list(self._output_structure, ret)
+        return [ops.convert_to_tensor(t) for t in ret]
+
+      func_name = defun_kwargs.pop("func_name", "wrapped_fn")
+      tf_function = def_function.Function(
+          python_function=wrapped_fn,
+          name=func_name,
+          input_signature=structure.get_flat_tensor_specs(
+              self._input_structure
+          ),
+          autograph=False,
+          experimental_attributes=defun_kwargs,
+      )
+
+      return tf_function.get_concrete_function
+
+    if use_legacy_function:
+      defun_kwargs.update({"func_name": func_name + "_" + str(ops.uid())})
+      fn_factory = trace_legacy_function(defun_kwargs)
+    else:
+      defun_kwargs.update({"func_name": func_name})
+      defun_kwargs.update({"_tf_data_function": True})
+      if debug_mode.DEBUG_MODE:
+        fn_factory = trace_py_function(defun_kwargs)
+      else:
+        if def_function.functions_run_eagerly():
+          warnings.warn(
+              "Even though the `tf.config.experimental_run_functions_eagerly` "
+              "option is set, this option does not apply to tf.data functions. "
+              "To force eager execution of tf.data functions, please use "
+              "`tf.data.experimental.enable_debug_mode()`.")
+        fn_factory = trace_tf_function(defun_kwargs)
+
+    self._function = fn_factory()
+    # There is no graph to add in eager mode.
+    add_to_graph &= not context.executing_eagerly()
+    # There are some lifetime issues when a legacy function is not added to a
+    # out-living graph. It's already deprecated so de-prioritizing the fix.
+    add_to_graph |= use_legacy_function
+    if add_to_graph:
+      self._function.add_to_graph(ops.get_default_graph())
+
+    if not use_legacy_function:
+      outer_graph_seed = ops.get_default_graph().seed
+      if outer_graph_seed and self._function.graph.seed == outer_graph_seed:
+        if self._function.graph._seed_used:
+          warnings.warn(
+              "Seed %s from outer graph might be getting used by function %s, "
+              "if the random op has not been provided any seed. Explicitly set "
+              "the seed in the function if this is not the intended behavior." %
+              (outer_graph_seed, func_name),
+              stacklevel=4)
+
+  @property
+  def output_structure(self):
+    return self._output_structure
+
+  @property
+  def output_classes(self):
+    return nest.map_structure(
+        lambda component_spec: component_spec._to_legacy_output_classes(),  # pylint: disable=protected-access
+        self._output_structure)
+
+  @property
+  def output_shapes(self):
+    return nest.map_structure(
+        lambda component_spec: component_spec._to_legacy_output_shapes(),  # pylint: disable=protected-access
+        self._output_structure)
+
+  @property
+  def output_types(self):
+    return nest.map_structure(
+        lambda component_spec: component_spec._to_legacy_output_types(),  # pylint: disable=protected-access
+        self._output_structure)
+
+  @property
+  def function(self):
+    return self._function

tensorflow/python/data/ops/take_op.py

@@ -0,0 +1,39 @@
+# 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 implementation of `tf.data.Dataset.skip`."""
+
+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
+
+
+def _take(self, count, name=None):  # pylint: disable=unused-private-name
+  return _TakeDataset(self, count, name=name)
+
+
+class _TakeDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A `Dataset` containing the first `count` elements from its input."""
+
+  def __init__(self, input_dataset, count, name=None):
+    """See `Dataset.take()` for details."""
+    self._input_dataset = input_dataset
+    self._count = ops.convert_to_tensor(count, dtype=dtypes.int64, name="count")
+    self._name = name
+    variant_tensor = gen_dataset_ops.take_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        count=self._count,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/take_while_op.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.take_while`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _take_while(input_dataset, predicate, name=None):  # pylint: disable=unused-private-name
+  """See `Dataset.take_while()` for details."""
+  return _TakeWhileDataset(input_dataset, predicate, name=name)
+
+
+class _TakeWhileDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A dataset that stops iteration when `predicate` returns false."""
+
+  def __init__(self, input_dataset, predicate, name=None):
+    """See `take_while()` for details."""
+
+    self._input_dataset = input_dataset
+    wrapped_func = structured_function.StructuredFunctionWrapper(
+        predicate, self._transformation_name(), dataset=self._input_dataset)
+
+    if not wrapped_func.output_structure.is_compatible_with(
+        tensor_spec.TensorSpec([], dtypes.bool)):
+      raise ValueError(f"Invalid `predicate`. `predicate` must return a "
+                       f"`tf.bool` scalar tensor but its return type is"
+                       f"{wrapped_func.output_structure}.")
+
+    self._predicate = wrapped_func
+    self._name = name
+    variant_tensor = ged_ops.take_while_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        other_arguments=self._predicate.function.captured_inputs,
+        predicate=self._predicate.function,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  def _functions(self):
+    return [self._predicate]
+
+  def _transformation_name(self):
+    return "Dataset.take_while()"

tensorflow/python/data/ops/test_mode.py

@@ -0,0 +1,32 @@
+# 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.
+# ==============================================================================
+"""Python test mode enabler.
+
+Enables test mode for tf.data.
+
+The test mode can be used to set up custom values for features and
+experiments as required in the unit tests.
+
+For example, if `warm_start` feature needs to be enabled exclusively for the
+unit tests, the tests can enable the test mode using `toggle_test_mode` and
+the default value of `warm_start` can be set as per the value of `TEST_MODE`.
+"""
+
+TEST_MODE = False
+
+
+def toggle_test_mode(test_mode):
+  global TEST_MODE
+  TEST_MODE = test_mode

tensorflow/python/data/ops/unbatch_op.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.unbatch`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+
+
+def _unbatch(input_dataset, name=None):  # pylint: disable=unused-private-name
+  """See `Dataset.unbatch()` for details."""
+  normalized_dataset = dataset_ops.normalize_to_dense(input_dataset)
+  return _UnbatchDataset(normalized_dataset, name=name)
+
+
+class _UnbatchDataset(dataset_ops.UnaryDataset):
+  """A dataset that splits the elements of its input into multiple elements."""
+
+  def __init__(self, input_dataset, name=None):
+    """See `unbatch()` for more details."""
+    flat_shapes = input_dataset._flat_shapes  # pylint: disable=protected-access
+    if any(s.ndims == 0 for s in flat_shapes):
+      raise ValueError("Cannot unbatch an input with scalar components.")
+    known_batch_dim = tensor_shape.Dimension(None)
+    for s in flat_shapes:
+      try:
+        known_batch_dim = known_batch_dim.merge_with(s[0])
+      except ValueError as e:
+        raise ValueError(
+            f"`unbatch()` is only supported for datasets of elements whose "
+            f"components have a matching leading dimension. Encountered both "
+            f"{known_batch_dim} and {s[0]}.") from e
+    self._input_dataset = input_dataset
+    self._structure = nest.map_structure(
+        lambda component_spec: component_spec._unbatch(),  # pylint: disable=protected-access
+        dataset_ops.get_structure(input_dataset))
+    self._name = name
+    variant_tensor = ged_ops.unbatch_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/unique_op.py

@@ -0,0 +1,42 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.unique`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_experimental_dataset_ops
+
+
+def _unique(input_dataset, name):  # pylint: disable=unused-private-name
+  return _UniqueDataset(input_dataset, name)
+
+
+class _UniqueDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A dataset containing the unique elements of an input dataset."""
+
+  def __init__(self, input_dataset, name=None):
+    """See `tf.data.Dataset.unique` for details."""
+    self._input_dataset = input_dataset
+    for ty in nest.flatten(dataset_ops.get_legacy_output_types(input_dataset)):
+      if ty not in (dtypes.int32, dtypes.int64, dtypes.string):
+        raise TypeError(
+            f"`tf.data.Dataset.unique` does not support type {ty} -- only "
+            f"`tf.int32`, `tf.int64`, and `tf.string` are supported.")
+    self._name = name
+    variant_tensor = gen_experimental_dataset_ops.unique_dataset(
+        self._input_dataset._variant_tensor,  # pylint: disable=protected-access
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)

tensorflow/python/data/ops/window_op.py

@@ -0,0 +1,76 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.window`."""
+
+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.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gen_dataset_ops
+
+
+def _window(input_dataset, size, shift, stride, drop_remainder, name):
+  if shift is None:
+    shift = size
+  return _WindowDataset(
+      input_dataset, size, shift, stride, drop_remainder, name=name)
+
+
+class _WindowDataset(dataset_ops.UnaryDataset):
+  """A dataset that creates window datasets from the input elements."""
+
+  def __init__(self,
+               input_dataset,
+               size,
+               shift,
+               stride,
+               drop_remainder,
+               name=None):
+    """See `window()` for more details."""
+    self._input_dataset = input_dataset
+    self._size = ops.convert_to_tensor(size, dtype=dtypes.int64, name="size")
+    self._shift = ops.convert_to_tensor(shift, dtype=dtypes.int64, name="shift")
+    self._stride = ops.convert_to_tensor(
+        stride, dtype=dtypes.int64, name="stride")
+    self._drop_remainder = ops.convert_to_tensor(
+        drop_remainder, dtype=dtypes.bool, name="drop_remainder")
+    self._structure = nest.pack_sequence_as(
+        dataset_ops.get_legacy_output_classes(input_dataset),
+        [
+            dataset_ops.DatasetSpec(  # pylint: disable=g-complex-comprehension
+                structure.convert_legacy_structure(output_type, output_shape,
+                                                   output_class))
+            for output_class, output_shape, output_type in zip(
+                nest.flatten(
+                    dataset_ops.get_legacy_output_classes(input_dataset)),
+                nest.flatten(
+                    dataset_ops.get_legacy_output_shapes(input_dataset)),
+                nest.flatten(
+                    dataset_ops.get_legacy_output_types(input_dataset)))
+        ])
+    self._name = name
+    variant_tensor = gen_dataset_ops.window_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        size=self._size,
+        shift=self._shift,
+        stride=self._stride,
+        drop_remainder=self._drop_remainder,
+        **self._common_args)
+    super().__init__(input_dataset, variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._structure

tensorflow/python/data/ops/zip_op.py

@@ -0,0 +1,62 @@
+# 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.
+# ==============================================================================
+"""The implementation of `tf.data.Dataset.zip`."""
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import nest
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.types import data as data_types
+
+
+def _zip(datasets, name):  # pylint: disable=redefined-builtin
+  return _ZipDataset(datasets, name)
+
+
+class _ZipDataset(dataset_ops.DatasetV2):
+  """A `Dataset` that zips its inputs together."""
+
+  def __init__(self, datasets, name=None):
+    """See `Dataset.zip()` for details."""
+    for ds in nest.flatten(datasets):
+      if not isinstance(ds, data_types.DatasetV2):
+        if isinstance(ds, list):
+          raise TypeError(
+              "Invalid input to `zip`. Inputs are expected to be (nested)"
+              " structures of `tf.data.Dataset` objects. Python `list` is"
+              " not supported and you should use `tuple` instead."
+          )
+        else:
+          raise TypeError(
+              "Invalid input to `zip`. Inputs are expected to be (nested)"
+              " structures of `tf.data.Dataset` objects but"
+              f" encountered object of type {type(ds)}."
+          )
+    self._datasets = datasets
+    self._structure = nest.pack_sequence_as(
+        self._datasets, [ds.element_spec for ds in nest.flatten(self._datasets)]
+    )
+    self._name = name
+    variant_tensor = gen_dataset_ops.zip_dataset(
+        [ds._variant_tensor for ds in nest.flatten(self._datasets)],
+        **self._common_args,
+    )
+    super().__init__(variant_tensor)
+
+  def _inputs(self):
+    return nest.flatten(self._datasets)
+
+  @property
+  def element_spec(self):
+    return self._structure