chore: import upstream snapshot with attribution

tensorflow/python/data/util/BUILD

@@ -0,0 +1,236 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "nest",
+    srcs = ["nest.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:nest_util",
+    ],
+)
+
+py_test(
+    name = "nest_test",
+    size = "small",
+    srcs = ["nest_test.py"],
+    strict_deps = True,
+    deps = [
+        ":nest",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "sparse",
+    srcs = ["sparse.py"],
+    strict_deps = True,
+    deps = [
+        ":nest",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:sparse_ops",
+    ],
+)
+
+py_test(
+    name = "sparse_test",
+    size = "small",
+    srcs = ["sparse_test.py"],
+    strict_deps = True,
+    deps = [
+        ":nest",
+        ":sparse",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "structure",
+    srcs = ["structure.py"],
+    strict_deps = True,
+    deps = [
+        ":nest",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:none_tensor",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/types:internal",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest_util",
+        "//tensorflow/python/util:tf_export",
+        "@pypi//wrapt",
+    ],
+)
+
+py_test(
+    name = "structure_test",
+    size = "small",
+    srcs = ["structure_test.py"],
+    strict_deps = True,
+    deps = [
+        ":nest",
+        ":structure",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/ops/ragged:ragged_tensor_value",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+        "@pypi//wrapt",
+    ],
+)
+
+py_library(
+    name = "options",
+    srcs = ["options.py"],
+    strict_deps = True,
+    deps = ["@absl_py//absl/logging"],
+)
+
+py_test(
+    name = "options_test",
+    size = "small",
+    srcs = ["options_test.py"],
+    strict_deps = True,
+    deps = [
+        ":options",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "convert",
+    srcs = ["convert.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+    ],
+)
+
+py_test(
+    name = "convert_test",
+    size = "small",
+    srcs = ["convert_test.py"],
+    strict_deps = True,
+    deps = [
+        ":convert",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "random_seed",
+    srcs = ["random_seed.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+    ],
+)
+
+py_test(
+    name = "random_seed_test",
+    size = "medium",
+    srcs = ["random_seed_test.py"],
+    strict_deps = True,
+    deps = [
+        ":random_seed",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "traverse",
+    srcs = ["traverse.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:dtypes",
+    ],
+)
+
+py_test(
+    name = "traverse_test",
+    size = "small",
+    srcs = ["traverse_test.py"],
+    strict_deps = True,
+    deps = [
+        ":traverse",
+        "//tensorflow/python/compat",
+        "//tensorflow/python/data/experimental/ops:dataset_ops",
+        "//tensorflow/python/data/kernel_tests:test_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/data/util/convert.py

@@ -0,0 +1,67 @@
+# 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.
+# ==============================================================================
+"""Helpers constructing Datasets."""
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+
+
+def optional_param_to_tensor(argument_name,
+                             argument_value,
+                             argument_default=0,
+                             argument_dtype=dtypes.int64):
+  if argument_value is not None:
+    return ops.convert_to_tensor(
+        argument_value, dtype=argument_dtype, name=argument_name)
+  else:
+    return constant_op.constant(
+        argument_default, dtype=argument_dtype, name=argument_name)
+
+
+def partial_shape_to_tensor(shape_like):
+  """Returns a `tf.Tensor` that represents the given shape.
+
+  Args:
+    shape_like: A value that can be converted to a `tf.TensorShape` or a
+      `tf.Tensor`.
+
+  Returns:
+    A 1-D `tf.Tensor` of `tf.int64` elements representing the given shape, where
+    `-1` is substituted for any unknown dimensions.
+  """
+  try:
+    # First attempt to convert the input to a shape, and return the
+    # "canonical" tensor representation, which uses `-1` in place of
+    # `None`.
+    shape_like = tensor_shape.as_shape(shape_like)
+    return ops.convert_to_tensor(
+        [dim if dim is not None else -1 for dim in shape_like.as_list()],
+        dtype=dtypes.int64)
+  except (TypeError, ValueError):
+    # The argument was not trivially convertible to a
+    # `tf.TensorShape`, so fall back on the conversion to tensor
+    # machinery.
+    ret = ops.convert_to_tensor(shape_like, preferred_dtype=dtypes.int64)
+    if ret.shape.dims is not None and len(ret.shape.dims) != 1:
+      raise ValueError("The given shape {} must be a 1-D tensor of `tf.int64` "
+                       "values, but the shape was {}.".format(
+                           shape_like, ret.shape))
+    if ret.dtype != dtypes.int64:
+      raise TypeError("The given shape {} must be a 1-D tensor of `tf.int64` "
+                      "values, but the element type was {}.".format(
+                          shape_like, ret.dtype.name))
+
+    return ret

tensorflow/python/data/util/convert_test.py

@@ -0,0 +1,160 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for utilities working with user input."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.util import convert
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class ConvertTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testInteger(self):
+    resp = convert.optional_param_to_tensor("foo", 3)
+    self.assertEqual(3, self.evaluate(resp))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIntegerDefault(self):
+    resp = convert.optional_param_to_tensor("foo", None)
+    self.assertEqual(0, self.evaluate(resp))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testStringDefault(self):
+    resp = convert.optional_param_to_tensor("bar", None, "default",
+                                            dtypes.string)
+    self.assertEqual(compat.as_bytes("default"), self.evaluate(resp))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testString(self):
+    resp = convert.optional_param_to_tensor("bar", "value", "default",
+                                            dtypes.string)
+    self.assertEqual(compat.as_bytes("value"), self.evaluate(resp))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPartialShapeToTensorKnownDimension(self):
+    self.assertAllEqual([1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([1]))))
+    self.assertAllEqual([1], self.evaluate(
+        convert.partial_shape_to_tensor((1,))))
+    self.assertAllEqual([1], self.evaluate(
+        convert.partial_shape_to_tensor([1])))
+    self.assertAllEqual([1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([1], dtype=dtypes.int64))))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testPartialShapeToTensorUnknownDimension(self):
+    self.assertAllEqual([-1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([None]))))
+    self.assertAllEqual([-1],
+                        self.evaluate(convert.partial_shape_to_tensor((None,))))
+    self.assertAllEqual([-1],
+                        self.evaluate(convert.partial_shape_to_tensor([None])))
+    self.assertAllEqual([-1],
+                        self.evaluate(convert.partial_shape_to_tensor([-1])))
+    self.assertAllEqual([-1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([-1],
+                                                     dtype=dtypes.int64))))
+
+    with self.assertRaisesRegex(
+        ValueError, r"The given shape .* must be a 1-D tensor of `tf.int64` "
+        r"values, but the shape was \(2, 2\)."):
+      convert.partial_shape_to_tensor(constant_op.constant(
+          [[1, 1], [1, 1]], dtype=dtypes.int64))
+
+    with self.assertRaisesRegex(
+        TypeError, r"The given shape .* must be a 1-D tensor of `tf.int64` "
+        r"values, but the element type was float32."):
+      convert.partial_shape_to_tensor(constant_op.constant([1., 1.]))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPartialShapeToTensorMultipleDimensions(self):
+    self.assertAllEqual([3, 6],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([3, 6]))))
+    self.assertAllEqual([3, 6],
+                        self.evaluate(convert.partial_shape_to_tensor((3, 6))))
+    self.assertAllEqual([3, 6],
+                        self.evaluate(convert.partial_shape_to_tensor([3, 6])))
+    self.assertAllEqual([3, 6],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([3, 6],
+                                                     dtype=dtypes.int64))))
+
+    self.assertAllEqual([3, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([3, None]))))
+    self.assertAllEqual([3, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor((3, None))))
+    self.assertAllEqual([3, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor([3, None])))
+    self.assertAllEqual([3, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([3, -1],
+                                                     dtype=dtypes.int64))))
+
+    self.assertAllEqual([-1, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([None, None]))))
+    self.assertAllEqual([-1, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor((None, None))))
+    self.assertAllEqual([-1, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor([None, None])))
+    self.assertAllEqual([-1, -1],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([-1, -1],
+                                                     dtype=dtypes.int64))))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPartialShapeToTensorScalar(self):
+    self.assertAllEqual([],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                tensor_shape.TensorShape([]))))
+    self.assertAllEqual([], self.evaluate(convert.partial_shape_to_tensor(())))
+    self.assertAllEqual([], self.evaluate(convert.partial_shape_to_tensor([])))
+    self.assertAllEqual([],
+                        self.evaluate(
+                            convert.partial_shape_to_tensor(
+                                constant_op.constant([], dtype=dtypes.int64))))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/util/nest.py

@@ -0,0 +1,305 @@
+# 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.
+# ==============================================================================
+
+"""## Functions for working with arbitrarily nested sequences of elements.
+
+NOTE(mrry): This fork of the `tensorflow.python.util.nest` module
+makes two changes:
+
+1. It removes support for lists as a level of nesting in nested structures.
+2. It adds support for `SparseTensorValue` as an atomic element.
+
+The motivation for this change is twofold:
+
+1. It seems more natural for lists to be treated (e.g. in Dataset constructors)
+   as tensors, rather than lists of (lists of...) tensors.
+2. This is needed because `SparseTensorValue` is implemented as a `namedtuple`
+   that would normally be flattened and we want to be able to create sparse
+   tensor from `SparseTensorValue's similarly to creating tensors from numpy
+   arrays.
+"""
+
+from tensorflow.python.util import nest_util
+
+
+def is_nested(structure):
+  return nest_util.is_nested(nest_util.Modality.DATA, structure)
+
+
+def flatten(structure):
+  return nest_util.flatten(nest_util.Modality.DATA, structure)
+
+
+def assert_same_structure(nest1, nest2, check_types=True):
+  """Asserts that two structures are nested in the same way.
+
+  Args:
+    nest1: an arbitrarily nested structure.
+    nest2: an arbitrarily nested structure.
+    check_types: if `True` (default) types of sequences should be same as
+      well. For dictionary, "type" of dictionary is considered to include its
+      keys. In other words, two dictionaries with different keys are considered
+      to have a different "type". If set to `False`, two iterables are
+      considered same as long as they yield the elements that have same
+      structures.
+
+  Raises:
+    ValueError: If the two structures do not have the same number of elements or
+      if the two structures are not nested in the same way.
+    TypeError: If the two structures differ in the type of sequence in any of
+      their substructures. Only possible if `check_types` is `True`.
+  """
+  nest_util.assert_same_structure(
+      nest_util.Modality.DATA, nest1, nest2, check_types
+  )
+
+
+def pack_sequence_as(structure, flat_sequence):
+  """Returns a given flattened sequence packed into a nest.
+
+  If `structure` is a scalar, `flat_sequence` must be a single-element list;
+  in this case the return value is `flat_sequence[0]`.
+
+  Args:
+    structure: tuple or list constructed of scalars and/or other tuples/lists,
+      or a scalar.  Note: numpy arrays are considered scalars.
+    flat_sequence: flat sequence to pack.
+
+  Returns:
+    packed: `flat_sequence` converted to have the same recursive structure as
+      `structure`.
+
+  Raises:
+    ValueError: If nest and structure have different element counts.
+  """
+  return nest_util.pack_sequence_as(
+      nest_util.Modality.DATA, structure, flat_sequence, expand_composites=False
+  )
+
+
+def map_structure(func, *structure, **check_types_dict):
+  """Applies `func` to each entry in `structure` and returns a new structure.
+
+  Applies `func(x[0], x[1], ...)` where x[i] is an entry in
+  `structure[i]`.  All structures in `structure` must have the same arity,
+  and the return value will contain the results in the same structure.
+
+  Args:
+    func: A callable that accepts as many arguments are there are structures.
+    *structure: scalar, or tuple or list of constructed scalars and/or other
+      tuples/lists, or scalars.  Note: numpy arrays are considered scalars.
+    **check_types_dict: only valid keyword argument is `check_types`. If set to
+      `True` (default) the types of iterables within the structures have to be
+      same (e.g. `map_structure(func, [1], (1,))` raises a `TypeError`
+      exception). To allow this set this argument to `False`.
+
+  Returns:
+    A new structure with the same arity as `structure`, whose values correspond
+    to `func(x[0], x[1], ...)` where `x[i]` is a value in the corresponding
+    location in `structure[i]`. If there are different sequence types and
+    `check_types` is `False` the sequence types of the first structure will be
+    used.
+
+  Raises:
+    TypeError: If `func` is not callable or if the structures do not match
+      each other by depth tree.
+    ValueError: If no structure is provided or if the structures do not match
+      each other by type.
+    ValueError: If wrong keyword arguments are provided.
+  """
+  return nest_util.map_structure(
+      nest_util.Modality.DATA, func, *structure, **check_types_dict
+  )
+
+
+def assert_shallow_structure(shallow_tree, input_tree, check_types=True):
+  """Asserts that `shallow_tree` is a shallow structure of `input_tree`.
+
+  That is, this function tests if the `input_tree` structure can be created from
+  the `shallow_tree` structure by replacing its leaf nodes with deeper
+  tree structures.
+
+  Examples:
+
+  The following code will raise an exception:
+  ```python
+    shallow_tree = ["a", "b"]
+    input_tree = ["c", ["d", "e"], "f"]
+    assert_shallow_structure(shallow_tree, input_tree)
+  ```
+
+  The following code will not raise an exception:
+  ```python
+    shallow_tree = ["a", "b"]
+    input_tree = ["c", ["d", "e"]]
+    assert_shallow_structure(shallow_tree, input_tree)
+  ```
+
+  Args:
+    shallow_tree: an arbitrarily nested structure.
+    input_tree: an arbitrarily nested structure.
+    check_types: if `True` (default) the sequence types of `shallow_tree` and
+      `input_tree` have to be the same.
+
+  Raises:
+    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
+    TypeError: If the sequence types of `shallow_tree` are different from
+      `input_tree`. Only raised if `check_types` is `True`.
+    ValueError: If the sequence lengths of `shallow_tree` are different from
+      `input_tree`.
+  """
+  nest_util.assert_shallow_structure(
+      nest_util.Modality.DATA, shallow_tree, input_tree, check_types
+  )
+
+
+def flatten_up_to(shallow_tree, input_tree):
+  """Flattens `input_tree` up to `shallow_tree`.
+
+  Any further depth in structure in `input_tree` is retained as elements in the
+  partially flatten output.
+
+  If `shallow_tree` and `input_tree` are not sequences, this returns a
+  single-element list: `[input_tree]`.
+
+  Use Case:
+
+  Sometimes we may wish to partially flatten a nested sequence, retaining some
+  of the nested structure. We achieve this by specifying a shallow structure,
+  `shallow_tree`, we wish to flatten up to.
+
+  The input, `input_tree`, can be thought of as having the same structure as
+  `shallow_tree`, but with leaf nodes that are themselves tree structures.
+
+  Examples:
+
+  ```python
+  input_tree = [[[2, 2], [3, 3]], [[4, 9], [5, 5]]]
+  shallow_tree = [[True, True], [False, True]]
+
+  flattened_input_tree = flatten_up_to(shallow_tree, input_tree)
+  flattened_shallow_tree = flatten_up_to(shallow_tree, shallow_tree)
+
+  # Output is:
+  # [[2, 2], [3, 3], [4, 9], [5, 5]]
+  # [True, True, False, True]
+  ```
+
+  ```python
+  input_tree = [[('a', 1), [('b', 2), [('c', 3), [('d', 4)]]]]]
+  shallow_tree = [['level_1', ['level_2', ['level_3', ['level_4']]]]]
+
+  input_tree_flattened_as_shallow_tree = flatten_up_to(shallow_tree, input_tree)
+  input_tree_flattened = flatten(input_tree)
+
+  # Output is:
+  # [('a', 1), ('b', 2), ('c', 3), ('d', 4)]
+  # ['a', 1, 'b', 2, 'c', 3, 'd', 4]
+  ```
+
+  Non-Sequence Edge Cases:
+
+  ```python
+  flatten_up_to(0, 0)  # Output: [0]
+  flatten_up_to(0, [0, 1, 2])  # Output: [[0, 1, 2]]
+  flatten_up_to([0, 1, 2], 0)  # Output: TypeError
+  flatten_up_to([0, 1, 2], [0, 1, 2])  # Output: [0, 1, 2]
+  ```
+
+  Args:
+    shallow_tree: a possibly pruned structure of input_tree.
+    input_tree: an arbitrarily nested structure or a scalar object.
+      Note, numpy arrays are considered scalars.
+
+  Returns:
+    A Python list, the partially flattened version of `input_tree` according to
+    the structure of `shallow_tree`.
+
+  Raises:
+    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
+    TypeError: If the sequence types of `shallow_tree` are different from
+      `input_tree`.
+    ValueError: If the sequence lengths of `shallow_tree` are different from
+      `input_tree`.
+  """
+  return nest_util.flatten_up_to(
+      nest_util.Modality.DATA, shallow_tree, input_tree
+  )
+
+
+def map_structure_up_to(shallow_tree, func, *inputs):
+  """Applies a function or op to a number of partially flattened inputs.
+
+  The `inputs` are flattened up to `shallow_tree` before being mapped.
+
+  Use Case:
+
+  Sometimes we wish to apply a function to a partially flattened
+  sequence (for example when the function itself takes sequence inputs). We
+  achieve this by specifying a shallow structure, `shallow_tree` we wish to
+  flatten up to.
+
+  The `inputs`, can be thought of as having the same structure as
+  `shallow_tree`, but with leaf nodes that are themselves tree structures.
+
+  This function, therefore, will return something with the same base structure
+  as `shallow_tree`.
+
+  Examples:
+
+  ```python
+  ab_tuple = collections.namedtuple("ab_tuple", "a, b")
+  op_tuple = collections.namedtuple("op_tuple", "add, mul")
+  inp_val = ab_tuple(a=2, b=3)
+  inp_ops = ab_tuple(a=op_tuple(add=1, mul=2), b=op_tuple(add=2, mul=3))
+  out = map_structure_up_to(inp_val, lambda val, ops: (val + ops.add) * ops.mul,
+                            inp_val, inp_ops)
+
+  # Output is: ab_tuple(a=6, b=15)
+  ```
+
+  ```python
+  data_list = [[2, 4, 6, 8], [[1, 3, 5, 7, 9], [3, 5, 7]]]
+  name_list = ['evens', ['odds', 'primes']]
+  out = map_structure_up_to(
+      name_list,
+      lambda name, sec: "first_{}_{}".format(len(sec), name),
+      name_list, data_list)
+
+  # Output is: ['first_4_evens', ['first_5_odds', 'first_3_primes']]
+  ```
+
+  Args:
+    shallow_tree: a shallow tree, common to all the inputs.
+    func: callable which will be applied to each input individually.
+    *inputs: arbitrarily nested combination of objects that are compatible with
+        shallow_tree. The function `func` is applied to corresponding
+        partially flattened elements of each input, so the function must support
+        arity of `len(inputs)`.
+
+  Raises:
+    TypeError: If `shallow_tree` is a sequence but `input_tree` is not.
+    TypeError: If the sequence types of `shallow_tree` are different from
+      `input_tree`.
+    ValueError: If the sequence lengths of `shallow_tree` are different from
+      `input_tree`.
+
+  Returns:
+    result of repeatedly applying `func`, with same structure as
+    `shallow_tree`.
+  """
+  return nest_util.map_structure_up_to(
+      nest_util.Modality.DATA, shallow_tree, func, *inputs
+  )

tensorflow/python/data/util/nest_test.py

@@ -0,0 +1,557 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for utilities working with arbitrarily nested structures."""
+
+import collections
+import dataclasses
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import test
+
+
+@dataclasses.dataclass
+class MaskedTensor:
+  mask: bool
+  value: tensor.Tensor
+
+  def __tf_flatten__(self):
+    metadata = (self.mask,)
+    components = (self.value,)
+    return metadata, components
+
+  def __tf_unflatten__(self, metadata, components):
+    mask = metadata[0]
+    value = components[0]
+    return MaskedTensor(mask=mask, value=value)
+
+
+class NestTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenAndPack(self):
+    structure = ((3, 4), 5, (6, 7, (9, 10), 8))
+    flat = ["a", "b", "c", "d", "e", "f", "g", "h"]
+    self.assertEqual(nest.flatten(structure), [3, 4, 5, 6, 7, 9, 10, 8])
+    self.assertEqual(
+        nest.pack_sequence_as(structure, flat), (("a", "b"), "c",
+                                                 ("d", "e", ("f", "g"), "h")))
+    point = collections.namedtuple("Point", ["x", "y"])
+    structure = (point(x=4, y=2), ((point(x=1, y=0),),))
+    flat = [4, 2, 1, 0]
+    self.assertEqual(nest.flatten(structure), flat)
+    restructured_from_flat = nest.pack_sequence_as(structure, flat)
+    self.assertEqual(restructured_from_flat, structure)
+    self.assertEqual(restructured_from_flat[0].x, 4)
+    self.assertEqual(restructured_from_flat[0].y, 2)
+    self.assertEqual(restructured_from_flat[1][0][0].x, 1)
+    self.assertEqual(restructured_from_flat[1][0][0].y, 0)
+
+    self.assertEqual([5], nest.flatten(5))
+    self.assertEqual([np.array([5])], nest.flatten(np.array([5])))
+
+    self.assertEqual("a", nest.pack_sequence_as(5, ["a"]))
+    self.assertEqual(
+        np.array([5]), nest.pack_sequence_as("scalar", [np.array([5])]))
+
+    with self.assertRaisesRegex(ValueError, "Argument `structure` is a scalar"):
+      nest.pack_sequence_as("scalar", [4, 5])
+
+    with self.assertRaisesRegex(TypeError, "flat_sequence"):
+      nest.pack_sequence_as([4, 5], "bad_sequence")
+
+    with self.assertRaises(ValueError):
+      nest.pack_sequence_as([5, 6, [7, 8]], ["a", "b", "c"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataclassIsNested(self):
+    mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    self.assertTrue(nest.is_nested(mt))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenDataclass(self):
+    mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    leaves = nest.flatten(mt)
+    self.assertLen(leaves, 1)
+    self.assertAllEqual(leaves[0], [1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPackDataclass(self):
+    mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    leaves = nest.flatten(mt)
+    reconstructed_mt = nest.pack_sequence_as(mt, leaves)
+    self.assertIsInstance(reconstructed_mt, MaskedTensor)
+    self.assertEqual(reconstructed_mt.mask, mt.mask)
+    self.assertAllEqual(reconstructed_mt.value, mt.value)
+
+    mt2 = MaskedTensor(mask=False, value=constant_op.constant([2]))
+    reconstructed_mt = nest.pack_sequence_as(mt2, leaves)
+    self.assertIsInstance(reconstructed_mt, MaskedTensor)
+    self.assertFalse(reconstructed_mt.mask)
+    self.assertAllEqual(reconstructed_mt.value, [1])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataclassMapStructure(self):
+    mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    mt_doubled = nest.map_structure(lambda x: x * 2, mt)
+    self.assertIsInstance(mt_doubled, MaskedTensor)
+    self.assertEqual(mt_doubled.mask, True)
+    self.assertAllEqual(mt_doubled.value, [2])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataclassAssertSameStructure(self):
+    mt1 = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    mt2 = MaskedTensor(mask=False, value=constant_op.constant([2]))
+    nest.assert_same_structure(mt1, mt2)
+
+    mt3 = (1, 2)
+
+    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
+        TypeError,
+        "don't have the same nested structure",
+    ):
+      nest.assert_same_structure(mt1, mt3)
+
+    class SubMaskedTensor(MaskedTensor):
+      pass
+
+    mt_subclass = SubMaskedTensor(mask=True, value=constant_op.constant([1]))
+    nest.assert_same_structure(mt1, mt_subclass, check_types=False)
+    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
+        TypeError,
+        "don't have the same sequence type",
+    ):
+      nest.assert_same_structure(mt1, mt_subclass)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDataclassAssertShallowStructure(self):
+    mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
+    structure1 = ("a", "b")
+    structure2 = (mt, "c")
+    nest.assert_shallow_structure(structure1, structure2)
+
+    structure3 = (mt, "d", "e")
+
+    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
+        ValueError,
+        "don't have the same sequence length",
+    ):
+      nest.assert_shallow_structure(structure1, structure3)
+
+    structure4 = {"a": mt, "b": "c"}
+    nest.assert_shallow_structure(structure1, structure4, check_types=False)
+    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
+        TypeError,
+        "don't have the same sequence type",
+    ):
+      nest.assert_shallow_structure(structure1, structure4)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenDictOrder(self):
+    """`flatten` orders dicts by key, including OrderedDicts."""
+    ordered = collections.OrderedDict([("d", 3), ("b", 1), ("a", 0), ("c", 2)])
+    plain = {"d": 3, "b": 1, "a": 0, "c": 2}
+    ordered_flat = nest.flatten(ordered)
+    plain_flat = nest.flatten(plain)
+    self.assertEqual([0, 1, 2, 3], ordered_flat)
+    self.assertEqual([0, 1, 2, 3], plain_flat)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testPackDictOrder(self):
+    """Packing orders dicts by key, including OrderedDicts."""
+    ordered = collections.OrderedDict([("d", 0), ("b", 0), ("a", 0), ("c", 0)])
+    plain = {"d": 0, "b": 0, "a": 0, "c": 0}
+    seq = [0, 1, 2, 3]
+    ordered_reconstruction = nest.pack_sequence_as(ordered, seq)
+    plain_reconstruction = nest.pack_sequence_as(plain, seq)
+    self.assertEqual(
+        collections.OrderedDict([("d", 3), ("b", 1), ("a", 0), ("c", 2)]),
+        ordered_reconstruction)
+    self.assertEqual({"d": 3, "b": 1, "a": 0, "c": 2}, plain_reconstruction)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenAndPackWithDicts(self):
+    # A nice messy mix of tuples, lists, dicts, and `OrderedDict`s.
+    named_tuple = collections.namedtuple("A", ("b", "c"))
+    mess = (
+        "z",
+        named_tuple(3, 4),
+        {
+            "c": (
+                1,
+                collections.OrderedDict([
+                    ("b", 3),
+                    ("a", 2),
+                ]),
+            ),
+            "b": 5
+        },
+        17
+    )
+
+    flattened = nest.flatten(mess)
+    self.assertEqual(flattened, ["z", 3, 4, 5, 1, 2, 3, 17])
+
+    structure_of_mess = (
+        14,
+        named_tuple("a", True),
+        {
+            "c": (
+                0,
+                collections.OrderedDict([
+                    ("b", 9),
+                    ("a", 8),
+                ]),
+            ),
+            "b": 3
+        },
+        "hi everybody",
+    )
+
+    unflattened = nest.pack_sequence_as(structure_of_mess, flattened)
+    self.assertEqual(unflattened, mess)
+
+    # Check also that the OrderedDict was created, with the correct key order.
+    unflattened_ordered_dict = unflattened[2]["c"][1]
+    self.assertIsInstance(unflattened_ordered_dict, collections.OrderedDict)
+    self.assertEqual(list(unflattened_ordered_dict.keys()), ["b", "a"])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenSparseValue(self):
+    st = sparse_tensor.SparseTensorValue([[0]], [0], [1])
+    single_value = st
+    list_of_values = [st, st, st]
+    nest_of_values = ((st), ((st), (st)))
+    dict_of_values = {"foo": st, "bar": st, "baz": st}
+    self.assertEqual([st], nest.flatten(single_value))
+    self.assertEqual([[st, st, st]], nest.flatten(list_of_values))
+    self.assertEqual([st, st, st], nest.flatten(nest_of_values))
+    self.assertEqual([st, st, st], nest.flatten(dict_of_values))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenRaggedValue(self):
+    rt = ragged_factory_ops.constant_value([[[0]], [[1]]])
+    single_value = rt
+    list_of_values = [rt, rt, rt]
+    nest_of_values = ((rt), ((rt), (rt)))
+    dict_of_values = {"foo": rt, "bar": rt, "baz": rt}
+    self.assertEqual([rt], nest.flatten(single_value))
+    self.assertEqual([[rt, rt, rt]], nest.flatten(list_of_values))
+    self.assertEqual([rt, rt, rt], nest.flatten(nest_of_values))
+    self.assertEqual([rt, rt, rt], nest.flatten(dict_of_values))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testIsNested(self):
+    self.assertFalse(nest.is_nested("1234"))
+    self.assertFalse(nest.is_nested([1, 3, [4, 5]]))
+    self.assertTrue(nest.is_nested(((7, 8), (5, 6))))
+    self.assertFalse(nest.is_nested([]))
+    self.assertFalse(nest.is_nested(set([1, 2])))
+    ones = array_ops.ones([2, 3])
+    self.assertFalse(nest.is_nested(ones))
+    self.assertFalse(nest.is_nested(math_ops.tanh(ones)))
+    self.assertFalse(nest.is_nested(np.ones((4, 5))))
+    self.assertTrue(nest.is_nested({"foo": 1, "bar": 2}))
+    self.assertFalse(
+        nest.is_nested(sparse_tensor.SparseTensorValue([[0]], [0], [1])))
+    self.assertFalse(
+        nest.is_nested(ragged_factory_ops.constant_value([[[0]], [[1]]])))
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertSameStructure(self):
+    structure1 = (((1, 2), 3), 4, (5, 6))
+    structure2 = ((("foo1", "foo2"), "foo3"), "foo4", ("foo5", "foo6"))
+    structure_different_num_elements = ("spam", "eggs")
+    structure_different_nesting = (((1, 2), 3), 4, 5, (6,))
+    structure_dictionary = {"foo": 2, "bar": 4, "baz": {"foo": 5, "bar": 6}}
+    structure_dictionary_diff_nested = {
+        "foo": 2,
+        "bar": 4,
+        "baz": {
+            "foo": 5,
+            "baz": 6
+        }
+    }
+    nest.assert_same_structure(structure1, structure2)
+    nest.assert_same_structure("abc", 1.0)
+    nest.assert_same_structure("abc", np.array([0, 1]))
+    nest.assert_same_structure("abc", constant_op.constant([0, 1]))
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure(structure1, structure_different_num_elements)
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure((0, 1), np.array([0, 1]))
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure(0, (0, 1))
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure(structure1, structure_different_nesting)
+
+    named_type_0 = collections.namedtuple("named_0", ("a", "b"))
+    named_type_1 = collections.namedtuple("named_1", ("a", "b"))
+    self.assertRaises(TypeError, nest.assert_same_structure, (0, 1),
+                      named_type_0("a", "b"))
+
+    nest.assert_same_structure(named_type_0(3, 4), named_type_0("a", "b"))
+
+    self.assertRaises(TypeError, nest.assert_same_structure,
+                      named_type_0(3, 4), named_type_1(3, 4))
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure(named_type_0(3, 4), named_type_0((3,), 4))
+
+    with self.assertRaisesRegex(ValueError,
+                                "don't have the same nested structure"):
+      nest.assert_same_structure(((3,), 4), (3, (4,)))
+
+    structure1_list = {"a": ((1, 2), 3), "b": 4, "c": (5, 6)}
+    structure2_list = {"a": ((1, 2), 3), "b": 4, "d": (5, 6)}
+    with self.assertRaisesRegex(TypeError, "don't have the same sequence type"):
+      nest.assert_same_structure(structure1, structure1_list)
+    nest.assert_same_structure(structure1, structure2, check_types=False)
+    nest.assert_same_structure(structure1, structure1_list, check_types=False)
+    with self.assertRaisesRegex(ValueError, "don't have the same set of keys"):
+      nest.assert_same_structure(structure1_list, structure2_list)
+    with self.assertRaisesRegex(ValueError, "don't have the same set of keys"):
+      nest.assert_same_structure(structure_dictionary,
+                                 structure_dictionary_diff_nested)
+    nest.assert_same_structure(
+        structure_dictionary,
+        structure_dictionary_diff_nested,
+        check_types=False)
+    nest.assert_same_structure(
+        structure1_list, structure2_list, check_types=False)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapStructure(self):
+    structure1 = (((1, 2), 3), 4, (5, 6))
+    structure2 = (((7, 8), 9), 10, (11, 12))
+    structure1_plus1 = nest.map_structure(lambda x: x + 1, structure1)
+    nest.assert_same_structure(structure1, structure1_plus1)
+    self.assertAllEqual(
+        [2, 3, 4, 5, 6, 7],
+        nest.flatten(structure1_plus1))
+    structure1_plus_structure2 = nest.map_structure(
+        lambda x, y: x + y, structure1, structure2)
+    self.assertEqual(
+        (((1 + 7, 2 + 8), 3 + 9), 4 + 10, (5 + 11, 6 + 12)),
+        structure1_plus_structure2)
+
+    self.assertEqual(3, nest.map_structure(lambda x: x - 1, 4))
+
+    self.assertEqual(7, nest.map_structure(lambda x, y: x + y, 3, 4))
+
+    with self.assertRaisesRegex(TypeError, "callable"):
+      nest.map_structure("bad", structure1_plus1)
+
+    with self.assertRaisesRegex(ValueError, "same nested structure"):
+      nest.map_structure(lambda x, y: None, 3, (3,))
+
+    with self.assertRaisesRegex(TypeError, "same sequence type"):
+      nest.map_structure(lambda x, y: None, ((3, 4), 5), {"a": (3, 4), "b": 5})
+
+    with self.assertRaisesRegex(ValueError, "same nested structure"):
+      nest.map_structure(lambda x, y: None, ((3, 4), 5), (3, (4, 5)))
+
+    with self.assertRaisesRegex(ValueError, "same nested structure"):
+      nest.map_structure(lambda x, y: None, ((3, 4), 5), (3, (4, 5)),
+                         check_types=False)
+
+    with self.assertRaisesRegex(ValueError, "Only valid keyword argument"):
+      nest.map_structure(lambda x: None, structure1, foo="a")
+
+    with self.assertRaisesRegex(ValueError, "Only valid keyword argument"):
+      nest.map_structure(lambda x: None, structure1, check_types=False, foo="a")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testAssertShallowStructure(self):
+    inp_ab = ("a", "b")
+    inp_abc = ("a", "b", "c")
+    expected_message = (
+        "The two structures don't have the same sequence length. Input "
+        "structure has length 2, while shallow structure has length 3.")
+    with self.assertRaisesRegex(ValueError, expected_message):
+      nest.assert_shallow_structure(inp_abc, inp_ab)
+
+    inp_ab1 = ((1, 1), (2, 2))
+    inp_ab2 = {"a": (1, 1), "b": (2, 2)}
+    expected_message = (
+        "The two structures don't have the same sequence type. Input structure "
+        "has type 'tuple', while shallow structure has type "
+        "'dict'.")
+    with self.assertRaisesRegex(TypeError, expected_message):
+      nest.assert_shallow_structure(inp_ab2, inp_ab1)
+    nest.assert_shallow_structure(inp_ab2, inp_ab1, check_types=False)
+
+    inp_ab1 = {"a": (1, 1), "b": {"c": (2, 2)}}
+    inp_ab2 = {"a": (1, 1), "b": {"d": (2, 2)}}
+    expected_message = (
+        r"The two structures don't have the same keys. Input "
+        r"structure has keys \['c'\], while shallow structure has "
+        r"keys \['d'\].")
+    with self.assertRaisesRegex(ValueError, expected_message):
+      nest.assert_shallow_structure(inp_ab2, inp_ab1)
+
+    inp_ab = collections.OrderedDict([("a", 1), ("b", (2, 3))])
+    inp_ba = collections.OrderedDict([("b", (2, 3)), ("a", 1)])
+    nest.assert_shallow_structure(inp_ab, inp_ba)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testFlattenUpTo(self):
+    input_tree = (((2, 2), (3, 3)), ((4, 9), (5, 5)))
+    shallow_tree = ((True, True), (False, True))
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [(2, 2), (3, 3), (4, 9), (5, 5)])
+    self.assertEqual(flattened_shallow_tree, [True, True, False, True])
+
+    input_tree = ((("a", 1), (("b", 2), (("c", 3), (("d", 4))))))
+    shallow_tree = (("level_1", ("level_2", ("level_3", ("level_4")))))
+    input_tree_flattened_as_shallow_tree = nest.flatten_up_to(shallow_tree,
+                                                              input_tree)
+    input_tree_flattened = nest.flatten(input_tree)
+    self.assertEqual(input_tree_flattened_as_shallow_tree,
+                     [("a", 1), ("b", 2), ("c", 3), ("d", 4)])
+    self.assertEqual(input_tree_flattened, ["a", 1, "b", 2, "c", 3, "d", 4])
+
+    ## Shallow non-list edge-case.
+    # Using iterable elements.
+    input_tree = ["input_tree"]
+    shallow_tree = "shallow_tree"
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    input_tree = ("input_tree_0", "input_tree_1")
+    shallow_tree = "shallow_tree"
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    # Using non-iterable elements.
+    input_tree = (0,)
+    shallow_tree = 9
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    input_tree = (0, 1)
+    shallow_tree = 9
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    ## Both non-list edge-case.
+    # Using iterable elements.
+    input_tree = "input_tree"
+    shallow_tree = "shallow_tree"
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    # Using non-iterable elements.
+    input_tree = 0
+    shallow_tree = 0
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [input_tree])
+    self.assertEqual(flattened_shallow_tree, [shallow_tree])
+
+    ## Input non-list edge-case.
+    # Using iterable elements.
+    input_tree = "input_tree"
+    shallow_tree = ("shallow_tree",)
+    expected_message = ("If shallow structure is a sequence, input must also "
+                        "be a sequence. Input has type: 'str'.")
+    with self.assertRaisesRegex(TypeError, expected_message):
+      flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_shallow_tree, list(shallow_tree))
+
+    input_tree = "input_tree"
+    shallow_tree = ("shallow_tree_9", "shallow_tree_8")
+    with self.assertRaisesRegex(TypeError, expected_message):
+      flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_shallow_tree, list(shallow_tree))
+
+    # Using non-iterable elements.
+    input_tree = 0
+    shallow_tree = (9,)
+    expected_message = ("If shallow structure is a sequence, input must also "
+                        "be a sequence. Input has type: 'int'.")
+    with self.assertRaisesRegex(TypeError, expected_message):
+      flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_shallow_tree, list(shallow_tree))
+
+    input_tree = 0
+    shallow_tree = (9, 8)
+    with self.assertRaisesRegex(TypeError, expected_message):
+      flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_shallow_tree, list(shallow_tree))
+
+    # Using dict.
+    input_tree = {"a": ((2, 2), (3, 3)), "b": ((4, 9), (5, 5))}
+    shallow_tree = {"a": (True, True), "b": (False, True)}
+    flattened_input_tree = nest.flatten_up_to(shallow_tree, input_tree)
+    flattened_shallow_tree = nest.flatten_up_to(shallow_tree, shallow_tree)
+    self.assertEqual(flattened_input_tree, [(2, 2), (3, 3), (4, 9), (5, 5)])
+    self.assertEqual(flattened_shallow_tree, [True, True, False, True])
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMapStructureUpTo(self):
+    ab_tuple = collections.namedtuple("ab_tuple", "a, b")
+    op_tuple = collections.namedtuple("op_tuple", "add, mul")
+    inp_val = ab_tuple(a=2, b=3)
+    inp_ops = ab_tuple(a=op_tuple(add=1, mul=2), b=op_tuple(add=2, mul=3))
+    out = nest.map_structure_up_to(
+        inp_val, lambda val, ops: (val + ops.add) * ops.mul, inp_val, inp_ops)
+    self.assertEqual(out.a, 6)
+    self.assertEqual(out.b, 15)
+
+    data_list = ((2, 4, 6, 8), ((1, 3, 5, 7, 9), (3, 5, 7)))
+    name_list = ("evens", ("odds", "primes"))
+    out = nest.map_structure_up_to(
+        name_list, lambda name, sec: "first_{}_{}".format(len(sec), name),
+        name_list, data_list)
+    self.assertEqual(out, ("first_4_evens", ("first_5_odds", "first_3_primes")))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/util/options.py

@@ -0,0 +1,177 @@
+# 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.
+# ==============================================================================
+"""Utilities for tf.data options."""
+
+import collections
+
+from absl import logging
+
+
+def _internal_attr_name(name):
+  return "_" + name
+
+
+class OptionsBase:
+  """Base class for representing a set of tf.data options.
+
+  Attributes:
+    _options: Stores the option values.
+  """
+
+  def __init__(self):
+    # NOTE: Cannot use `self._options` here as we override `__setattr__`
+    object.__setattr__(self, "_options", {})
+    object.__setattr__(self, "_mutable", True)
+
+  def __eq__(self, other):
+    if not isinstance(other, self.__class__):
+      return NotImplemented
+    for name in set(self._options) | set(other._options):  # pylint: disable=protected-access
+      if getattr(self, name) != getattr(other, name):
+        return False
+    return True
+
+  def __ne__(self, other):
+    if isinstance(other, self.__class__):
+      return not self.__eq__(other)
+    else:
+      return NotImplemented
+
+  def __setattr__(self, name, value):
+    if not self._mutable:
+      raise ValueError("Mutating `tf.data.Options()` returned by "
+                       "`tf.data.Dataset.options()` has no effect. Use "
+                       "`tf.data.Dataset.with_options(options)` to set or "
+                       "update dataset options.")
+    if hasattr(self, name):
+      object.__setattr__(self, name, value)
+    else:
+      raise AttributeError("Cannot set the property {} on {}.".format(
+          name,
+          type(self).__name__))
+
+  def _set_mutable(self, mutable):
+    """Change the mutability property to `mutable`."""
+    object.__setattr__(self, "_mutable", mutable)
+
+  def _to_proto(self):
+    """Convert options to protocol buffer."""
+    raise NotImplementedError("{}._to_proto()".format(type(self).__name__))
+
+  def _from_proto(self, pb):
+    """Convert protocol buffer to options."""
+    raise NotImplementedError("{}._from_proto()".format(type(self).__name__))
+
+
+# Creates a namedtuple with three keys for optimization graph rewrites settings.
+def graph_rewrites():
+  return collections.namedtuple("GraphRewrites",
+                                ["enabled", "disabled", "default"])
+
+
+def create_option(name, ty, docstring, default_factory=lambda: None):
+  """Creates a type-checked property.
+
+  Args:
+    name: The name to use.
+    ty: The type to use. The type of the property will be validated when it
+      is set.
+    docstring: The docstring to use.
+    default_factory: A callable that takes no arguments and returns a default
+      value to use if not set.
+
+  Returns:
+    A type-checked property.
+  """
+
+  def get_fn(option):
+    # pylint: disable=protected-access
+    if name not in option._options:
+      option._options[name] = default_factory()
+    return option._options.get(name)
+
+  def set_fn(option, value):
+    if not isinstance(value, ty):
+      raise TypeError(
+          "Property \"{}\" must be of type {}, got: {} (type: {})".format(
+              name, ty, value, type(value)))
+    option._options[name] = value  # pylint: disable=protected-access
+
+  return property(get_fn, set_fn, None, docstring)
+
+
+def merge_options(*options_list):
+  """Merges the given options, returning the result as a new options object.
+
+  The input arguments are expected to have a matching type that derives from
+  `tf.data.OptionsBase` (and thus each represent a set of options). The method
+  outputs an object of the same type created by merging the sets of options
+  represented by the input arguments.
+
+  If an option is set to different values by different options objects, the
+  result will match the setting of the options object that appears in the input
+  list last.
+
+  If an option is an instance of `tf.data.OptionsBase` itself, then this method
+  is applied recursively to the set of options represented by this option.
+
+  Args:
+    *options_list: options to merge
+
+  Raises:
+    TypeError: if the input arguments are incompatible or not derived from
+      `tf.data.OptionsBase`
+
+  Returns:
+    A new options object which is the result of merging the given options.
+  """
+  if len(options_list) < 1:
+    raise ValueError("At least one options should be provided")
+  result_type = type(options_list[0])
+
+  for options in options_list:
+    if not isinstance(options, result_type):
+      raise TypeError(
+          "Could not merge incompatible options of type {} and {}.".format(
+              type(options), result_type))
+
+  if not isinstance(options_list[0], OptionsBase):
+    raise TypeError(
+        "All options to be merged should inherit from `OptionsBase` but found "
+        "option of type {} which does not.".format(type(options_list[0])))
+
+  default_options = result_type()
+  result = result_type()
+  for options in options_list:
+    # Iterate over all set options and merge them into the result.
+    for name in options._options:  # pylint: disable=protected-access
+      this = getattr(result, name)
+      that = getattr(options, name)
+      default = getattr(default_options, name)
+      if that == default:
+        continue
+      elif this == default:
+        setattr(result, name, that)
+      elif isinstance(this, OptionsBase):
+        setattr(result, name, merge_options(this, that))
+      elif name == "framework_type":
+        # Since, `framework_type`` is a repeated string field (list), the merged
+        # result will be a combined list.
+        setattr(result, name, this+that)
+      elif this != that:
+        logging.warning("Changing the value of option %s from %r to %r.", name,
+                        this, that)
+        setattr(result, name, that)
+  return result

tensorflow/python/data/util/options_test.py

@@ -0,0 +1,141 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for dataset options utilities."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.util import options
+from tensorflow.python.framework import combinations
+from tensorflow.python.platform import test
+
+
+class _TestOptions(options.OptionsBase):
+  x = options.create_option(
+      name="x",
+      ty=int,
+      docstring="the answer to everything",
+      default_factory=lambda: 42)
+  y = options.create_option(
+      name="y", ty=float, docstring="a tasty pie", default_factory=lambda: 3.14)
+
+
+class _NestedTestOptions(options.OptionsBase):
+  opts = options.create_option(
+      name="opts", ty=_TestOptions, docstring="nested options")
+
+
+class OptionsTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testDocumentation(self):
+    self.assertEqual(_TestOptions.x.__doc__, "the answer to everything")
+    self.assertEqual(_TestOptions.y.__doc__, "a tasty pie")
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testCreateOption(self):
+    opts = _TestOptions()
+    self.assertEqual(opts.x, 42)
+    self.assertEqual(opts.y, 3.14)
+    self.assertIsInstance(opts.x, int)
+    self.assertIsInstance(opts.y, float)
+    opts.x = 0
+    self.assertEqual(opts.x, 0)
+    with self.assertRaises(TypeError):
+      opts.x = 3.14
+    opts.y = 0.0
+    self.assertEqual(opts.y, 0.0)
+    with self.assertRaises(TypeError):
+      opts.y = 42
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMergeOptions(self):
+    options1, options2 = _TestOptions(), _TestOptions()
+    with self.assertRaises(ValueError):
+      options.merge_options()
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.x, 42)
+    self.assertEqual(merged_options.y, 3.14)
+    options1.x = 0
+    options2.y = 0.0
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.x, 0)
+    self.assertEqual(merged_options.y, 0.0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMergeNestedOptions(self):
+    options1, options2 = _NestedTestOptions(), _NestedTestOptions()
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.opts, None)
+    options1.opts = _TestOptions()
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.opts, _TestOptions())
+    options2.opts = _TestOptions()
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.opts, _TestOptions())
+    options1.opts.x = 0
+    options2.opts.y = 0.0
+    merged_options = options.merge_options(options1, options2)
+    self.assertEqual(merged_options.opts.x, 0)
+    self.assertEqual(merged_options.opts.y, 0.0)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testImmutable(self):
+    test_options = _TestOptions()
+    test_options._set_mutable(False)
+
+    with self.assertRaisesRegex(
+        ValueError, r"Mutating `tf.data.Options\(\)` returned by "
+        r"`tf.data.Dataset.options\(\)` has no effect. Use "
+        r"`tf.data.Dataset.with_options\(options\)` to set or "
+        "update dataset options."):
+      test_options.test = 100
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testNoSpuriousAttrs(self):
+    test_options = _TestOptions()
+
+    with self.assertRaisesRegex(
+        AttributeError, "Cannot set the property wrong_attr on _TestOptions."):
+      test_options.wrong_attr = True
+    with self.assertRaises(AttributeError):
+      _ = test_options.wrong_attr
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMergeNoOptions(self):
+    with self.assertRaisesRegex(ValueError,
+                                "At least one options should be provided"):
+      options.merge_options()
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMergeOptionsDifferentType(self):
+    options1, options2 = _TestOptions(), _NestedTestOptions()
+    with self.assertRaisesRegex(
+        TypeError, r"Could not merge incompatible options of type "
+        r"\<class \'__main__._NestedTestOptions\'\> and "
+        r"\<class \'__main__._TestOptions\'\>."):
+      options.merge_options(options1, options2)
+
+  @combinations.generate(test_base.default_test_combinations())
+  def testMergeOptionsWrongType(self):
+    with self.assertRaisesRegex(
+        TypeError, "All options to be merged should inherit from "
+        r"\`OptionsBase\` but found option of type \<class \'int\'\> which "
+        "does not."):
+      options.merge_options(1, 2, 3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/util/random_seed.py

@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+"""Utilities for generating Tensor-valued random seeds."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+
+
+def get_seed(seed):
+  """Returns the local seeds an operation should use given an op-specific seed.
+
+  See `random_seed.get_seed` for more details. This wrapper adds support for
+  the case where `seed` may be a tensor.
+
+  Args:
+    seed: An integer or a `tf.int64` scalar tensor.
+
+  Returns:
+    A tuple of two `tf.int64` scalar tensors that should be used for the local
+    seed of the calling dataset.
+  """
+  seed, seed2 = random_seed.get_seed(seed)
+  if seed is None:
+    seed = constant_op.constant(0, dtype=dtypes.int64, name="seed")
+  else:
+    seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed")
+  if seed2 is None:
+    seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2")
+  else:
+    with ops.name_scope("seed2") as scope:
+      seed2 = ops.convert_to_tensor(seed2, dtype=dtypes.int64)
+      seed2 = array_ops.where_v2(
+          math_ops.logical_and(
+              math_ops.equal(seed, 0), math_ops.equal(seed2, 0)),
+          constant_op.constant(2**31 - 1, dtype=dtypes.int64),
+          seed2,
+          name=scope)
+  return seed, seed2

tensorflow/python/data/util/random_seed_test.py

@@ -0,0 +1,133 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for utilities working with arbitrarily nested structures."""
+
+import functools
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.util import random_seed as data_random_seed
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import random_seed
+from tensorflow.python.platform import test
+
+
+# NOTE(vikoth18): Arguments of parameterized tests are lifted into lambdas to make
+# sure they are not executed before the (eager- or graph-mode) test environment
+# has been set up.
+
+
+def _test_random_seed_combinations():
+
+  cases = [
+      # Each test case is a tuple with input to get_seed:
+      # (input_graph_seed, input_op_seed)
+      # and output from get_seed:
+      # (output_graph_seed, output_op_seed)
+      (
+          "TestCase_0",
+          lambda: (None, None),
+          lambda: (0, 0),
+      ),
+      ("TestCase_1", lambda: (None, 1), lambda:
+       (random_seed.DEFAULT_GRAPH_SEED, 1)),
+      ("TestCase_2", lambda: (1, 1), lambda: (1, 1)),
+      (
+          # Avoid nondeterministic (0, 0) output
+          "TestCase_3",
+          lambda: (0, 0),
+          lambda: (0, 2**31 - 1)),
+      (
+          # Don't wrap to (0, 0) either
+          "TestCase_4",
+          lambda: (2**31 - 1, 0),
+          lambda: (0, 2**31 - 1)),
+      (
+          # Wrapping for the other argument
+          "TestCase_5",
+          lambda: (0, 2**31 - 1),
+          lambda: (0, 2**31 - 1)),
+      (
+          # Once more, with tensor-valued arguments
+          "TestCase_6",
+          lambda:
+          (None, constant_op.constant(1, dtype=dtypes.int64, name="one")),
+          lambda: (random_seed.DEFAULT_GRAPH_SEED, 1)),
+      ("TestCase_7", lambda:
+       (1, constant_op.constant(1, dtype=dtypes.int64, name="one")), lambda:
+       (1, 1)),
+      (
+          "TestCase_8",
+          lambda: (0, constant_op.constant(0, dtype=dtypes.int64, name="zero")),
+          lambda: (0, 2**31 - 1)  # Avoid nondeterministic (0, 0) output
+      ),
+      (
+          "TestCase_9",
+          lambda:
+          (2**31 - 1, constant_op.constant(0, dtype=dtypes.int64, name="zero")),
+          lambda: (0, 2**31 - 1)  # Don't wrap to (0, 0) either
+      ),
+      (
+          "TestCase_10",
+          lambda:
+          (0, constant_op.constant(
+              2**31 - 1, dtype=dtypes.int64, name="intmax")),
+          lambda: (0, 2**31 - 1)  # Wrapping for the other argument
+      )
+  ]
+
+  def reduce_fn(x, y):
+    name, input_fn, output_fn = y
+    return x + combinations.combine(
+        input_fn=combinations.NamedObject("input_fn.{}".format(name), input_fn),
+        output_fn=combinations.NamedObject("output_fn.{}".format(name),
+                                           output_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class RandomSeedTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  def _checkEqual(self, tinput, toutput):
+    random_seed.set_random_seed(tinput[0])
+    g_seed, op_seed = data_random_seed.get_seed(tinput[1])
+    g_seed = self.evaluate(g_seed)
+    op_seed = self.evaluate(op_seed)
+    msg = "test_case = {0}, got {1}, want {2}".format(tinput, (g_seed, op_seed),
+                                                      toutput)
+    self.assertEqual((g_seed, op_seed), toutput, msg=msg)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_random_seed_combinations()))
+  def testRandomSeed(self, input_fn, output_fn):
+    tinput, toutput = input_fn(), output_fn()
+    self._checkEqual(tinput=tinput, toutput=toutput)
+    random_seed.set_random_seed(None)
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testIncrementalRandomSeed(self):
+    random_seed.set_random_seed(1)
+    for i in range(10):
+      tinput = (1, None)
+      toutput = (1, i)
+      self._checkEqual(tinput=tinput, toutput=toutput)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/data/util/sparse.py

@@ -0,0 +1,148 @@
+# 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 dataset sparse tensor utility functions."""
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import sparse_ops
+
+
+def any_sparse(classes):
+  """Checks for sparse tensor.
+
+  Args:
+    classes: a structure of objects that identify the dataset item classes
+
+  Returns:
+    `True` if `classes` contains a sparse tensor type and `False` otherwise.
+  """
+  return any(c is sparse_tensor.SparseTensor for c in nest.flatten(classes))
+
+
+def as_dense_shapes(shapes, classes):
+  """Converts sparse tensor shapes to their physical shapes.
+
+  Args:
+    shapes: a structure of shapes to convert.
+    classes: a structure of objects that identify the dataset item classes
+
+  Returns:
+    a structure matching the nested structure of `shapes`, containing
+    `tensor_shape.unknown_shape()` at positions where `classes` contains
+    `tf.sparse.SparseTensor` and matching contents of `shapes` otherwise
+  """
+  ret = nest.pack_sequence_as(shapes, [
+      tensor_shape.unknown_shape() if c is sparse_tensor.SparseTensor else shape
+      for shape, c in zip(nest.flatten(shapes), nest.flatten(classes))
+  ])
+  return ret
+
+
+def as_dense_types(types, classes):
+  """Converts sparse tensor types to `dtypes.variant`.
+
+  Args:
+    types: a structure of types to convert.
+    classes: a structure of objects that identify the dataset item classes
+
+  Returns:
+    a structure matching the nested structure of `types`, containing
+    `dtypes.variant` at positions where `classes` contains
+    `tf.sparse.SparseTensor` and matching contents of `types` otherwise
+  """
+  ret = nest.pack_sequence_as(types, [
+      dtypes.variant if c is sparse_tensor.SparseTensor else ty
+      for ty, c in zip(nest.flatten(types), nest.flatten(classes))
+  ])
+  return ret
+
+
+def deserialize_sparse_tensors(tensors, types, shapes, classes):
+  """Deserializes sparse tensors.
+
+  Args:
+    tensors: a structure of tensors to deserialize.
+    types: a structure that holds information about types of `tensors`
+    shapes: a structure that holds information about shapes of `tensors`
+    classes: a structure of objects that identify the dataset item classes
+
+  Returns:
+    `tensors` with any serialized sparse tensors replaced by their deserialized
+    version.
+  """
+  ret = nest.pack_sequence_as(types, [
+      sparse_ops.deserialize_sparse(tensor, dtype=ty, rank=shape.ndims)
+      if c is sparse_tensor.SparseTensor else tensor
+      for (tensor, ty, shape, c) in zip(
+          nest.flatten(tensors), nest.flatten(types), nest.flatten(shapes),
+          nest.flatten(classes))
+  ])
+  return ret
+
+
+def get_classes(tensors):
+  """Gets classes for a structure of tensors.
+
+  Args:
+    tensors: the tensor structure to get classes for.
+
+  Returns:
+    a structure matching the nested structure of `tensors`, containing
+    `tf.sparse.SparseTensor` at positions where `tensors` contains a sparse
+    tensor and `tf.Tensor` otherwise.
+  """
+  return nest.pack_sequence_as(tensors, [
+      sparse_tensor.SparseTensor
+      if isinstance(tensor, sparse_tensor.SparseTensor) else tensor_lib.Tensor
+      for tensor in nest.flatten(tensors)
+  ])
+
+
+def serialize_many_sparse_tensors(tensors):
+  """Serializes many sparse tensors into a batch.
+
+  Args:
+    tensors: a tensor structure to serialize.
+
+  Returns:
+    `tensors` with any sparse tensors replaced by the serialized batch.
+  """
+
+  ret = nest.pack_sequence_as(tensors, [
+      sparse_ops.serialize_many_sparse(tensor, out_type=dtypes.variant)
+      if sparse_tensor.is_sparse(tensor) else tensor
+      for tensor in nest.flatten(tensors)
+  ])
+  return ret
+
+
+def serialize_sparse_tensors(tensors):
+  """Serializes sparse tensors.
+
+  Args:
+    tensors: a tensor structure to serialize.
+
+  Returns:
+    `tensors` with any sparse tensors replaced by their serialized version.
+  """
+
+  ret = nest.pack_sequence_as(tensors, [
+      sparse_ops.serialize_sparse(tensor, out_type=dtypes.variant)
+      if isinstance(tensor, sparse_tensor.SparseTensor) else tensor
+      for tensor in nest.flatten(tensors)
+  ])
+  return ret

tensorflow/python/data/util/sparse_test.py

@@ -0,0 +1,346 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for utilities working with arbitrarily nested structures."""
+
+import functools
+
+from absl.testing import parameterized
+
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.util import nest
+from tensorflow.python.data.util import sparse
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+
+
+# NOTE(vikoth18): Arguments of parameterized tests are lifted into lambdas to make
+# sure they are not executed before the (eager- or graph-mode) test environment
+# has been set up.
+#
+
+
+def _test_any_sparse_combinations():
+
+  cases = [("TestCase_0", lambda: (), False),
+           ("TestCase_1", lambda: (tensor.Tensor), False),
+           ("TestCase_2", lambda: (((tensor.Tensor))), False),
+           ("TestCase_3", lambda: (tensor.Tensor, tensor.Tensor), False),
+           ("TestCase_4", lambda:
+            (tensor.Tensor, sparse_tensor.SparseTensor), True),
+           ("TestCase_5", lambda:
+            (sparse_tensor.SparseTensor, sparse_tensor.SparseTensor), True),
+           ("TestCase_6", lambda: (((sparse_tensor.SparseTensor))), True)]
+
+  def reduce_fn(x, y):
+    name, classes_fn, expected = y
+    return x + combinations.combine(
+        classes_fn=combinations.NamedObject("classes_fn.{}".format(name),
+                                            classes_fn),
+        expected=expected)
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _test_as_dense_shapes_combinations():
+
+  cases = [
+      ("TestCase_0", lambda: (), lambda: (), lambda: ()),
+      ("TestCase_1", lambda: tensor_shape.TensorShape([]),
+       lambda: tensor.Tensor,
+       lambda: tensor_shape.TensorShape([])),
+      (
+          "TestCase_2",
+          lambda: tensor_shape.TensorShape([]),
+          lambda: sparse_tensor.SparseTensor,
+          lambda: tensor_shape.unknown_shape()  # pylint: disable=unnecessary-lambda
+      ),
+      ("TestCase_3", lambda: (tensor_shape.TensorShape([])), lambda:
+       (tensor.Tensor), lambda: (tensor_shape.TensorShape([]))),
+      (
+          "TestCase_4",
+          lambda: (tensor_shape.TensorShape([])),
+          lambda: (sparse_tensor.SparseTensor),
+          lambda: (tensor_shape.unknown_shape())  # pylint: disable=unnecessary-lambda
+      ),
+      ("TestCase_5", lambda: (tensor_shape.TensorShape([]), ()), lambda:
+       (tensor.Tensor, ()), lambda: (tensor_shape.TensorShape([]), ())),
+      ("TestCase_6", lambda: ((), tensor_shape.TensorShape([])), lambda:
+       ((), tensor.Tensor), lambda: ((), tensor_shape.TensorShape([]))),
+      ("TestCase_7", lambda: (tensor_shape.TensorShape([]), ()), lambda:
+       (sparse_tensor.SparseTensor, ()), lambda: (tensor_shape.unknown_shape(),
+                                                  ())),
+      ("TestCase_8", lambda: ((), tensor_shape.TensorShape([])), lambda:
+       ((), sparse_tensor.SparseTensor), lambda: (
+           (), tensor_shape.unknown_shape())),
+      ("TestCase_9", lambda: (tensor_shape.TensorShape([]),
+                              (), tensor_shape.TensorShape([])), lambda:
+       (tensor.Tensor, (), tensor.Tensor), lambda:
+       (tensor_shape.TensorShape([]), (), tensor_shape.TensorShape([]))),
+      ("TestCase_10", lambda: (tensor_shape.TensorShape([]),
+                               (), tensor_shape.TensorShape([])), lambda:
+       (sparse_tensor.SparseTensor, (), sparse_tensor.SparseTensor), lambda:
+       (tensor_shape.unknown_shape(), (), tensor_shape.unknown_shape())),
+      ("TestCase_11", lambda: ((), tensor_shape.TensorShape([]), ()), lambda:
+       ((), tensor.Tensor, ()), lambda: ((), tensor_shape.TensorShape([]), ())),
+      ("TestCase_12", lambda: ((), tensor_shape.TensorShape([]), ()), lambda:
+       ((), sparse_tensor.SparseTensor,
+        ()), lambda: ((), tensor_shape.unknown_shape(), ()))
+  ]
+
+  def reduce_fn(x, y):
+    name, types_fn, classes_fn, expected_fn = y
+    return x + combinations.combine(
+        types_fn=combinations.NamedObject("types_fn.{}".format(name), types_fn),
+        classes_fn=combinations.NamedObject("classes_fn.{}".format(name),
+                                            classes_fn),
+        expected_fn=combinations.NamedObject("expected_fn.{}".format(name),
+                                             expected_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _test_as_dense_types_combinations():
+  cases = [
+      ("TestCase_0", lambda: (), lambda: (), lambda: ()),
+      ("TestCase_1", lambda: dtypes.int32, lambda: tensor.Tensor,
+       lambda: dtypes.int32),
+      ("TestCase_2", lambda: dtypes.int32, lambda: sparse_tensor.SparseTensor,
+       lambda: dtypes.variant),
+      ("TestCase_3", lambda: (dtypes.int32), lambda: (tensor.Tensor), lambda:
+       (dtypes.int32)),
+      ("TestCase_4", lambda: (dtypes.int32), lambda:
+       (sparse_tensor.SparseTensor), lambda: (dtypes.variant)),
+      ("TestCase_5", lambda: (dtypes.int32, ()), lambda:
+       (tensor.Tensor, ()), lambda: (dtypes.int32, ())),
+      ("TestCase_6", lambda: ((), dtypes.int32), lambda:
+       ((), tensor.Tensor), lambda: ((), dtypes.int32)),
+      ("TestCase_7", lambda: (dtypes.int32, ()), lambda:
+       (sparse_tensor.SparseTensor, ()), lambda: (dtypes.variant, ())),
+      ("TestCase_8", lambda: ((), dtypes.int32), lambda:
+       ((), sparse_tensor.SparseTensor), lambda: ((), dtypes.variant)),
+      ("TestCase_9", lambda: (dtypes.int32, (), dtypes.int32), lambda:
+       (tensor.Tensor, (), tensor.Tensor),
+       lambda: (dtypes.int32, (), dtypes.int32)),
+      ("TestCase_10", lambda: (dtypes.int32, (), dtypes.int32), lambda:
+       (sparse_tensor.SparseTensor, (), sparse_tensor.SparseTensor), lambda:
+       (dtypes.variant, (), dtypes.variant)),
+      ("TestCase_11", lambda: ((), dtypes.int32, ()), lambda:
+       ((), tensor.Tensor, ()), lambda: ((), dtypes.int32, ())),
+      ("TestCase_12", lambda: ((), dtypes.int32, ()), lambda:
+       ((), sparse_tensor.SparseTensor, ()), lambda: ((), dtypes.variant, ())),
+  ]
+
+  def reduce_fn(x, y):
+    name, types_fn, classes_fn, expected_fn = y
+    return x + combinations.combine(
+        types_fn=combinations.NamedObject("types_fn.{}".format(name), types_fn),
+        classes_fn=combinations.NamedObject("classes_fn.{}".format(name),
+                                            classes_fn),
+        expected_fn=combinations.NamedObject("expected_fn.{}".format(name),
+                                             expected_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _test_get_classes_combinations():
+  cases = [
+      ("TestCase_0", lambda: (), lambda: ()),
+      ("TestCase_1", lambda: sparse_tensor.SparseTensor(
+          indices=[[0]], values=[1], dense_shape=[1]),
+       lambda: sparse_tensor.SparseTensor),
+      ("TestCase_2", lambda: constant_op.constant([1]), lambda: tensor.Tensor),
+      ("TestCase_3", lambda:
+       (sparse_tensor.SparseTensor(indices=[[0]], values=[1], dense_shape=[1])),
+       lambda: (sparse_tensor.SparseTensor)),
+      ("TestCase_4", lambda: (constant_op.constant([1])),
+       lambda: (tensor.Tensor)),
+      ("TestCase_5", lambda:
+       (sparse_tensor.SparseTensor(indices=[[0]], values=[1], dense_shape=[1]),
+        ()), lambda: (sparse_tensor.SparseTensor, ())),
+      ("TestCase_6", lambda:
+       ((),
+        sparse_tensor.SparseTensor(indices=[[0]], values=[1], dense_shape=[1])),
+       lambda: ((), sparse_tensor.SparseTensor)),
+      ("TestCase_7", lambda: (constant_op.constant([1]), ()), lambda:
+       (tensor.Tensor, ())),
+      ("TestCase_8", lambda: ((), constant_op.constant([1])), lambda:
+       ((), tensor.Tensor)),
+      ("TestCase_9", lambda:
+       (sparse_tensor.SparseTensor(indices=[[0]], values=[1], dense_shape=[1]),
+        (), constant_op.constant([1])), lambda: (sparse_tensor.SparseTensor,
+                                                 (), tensor.Tensor)),
+      ("TestCase_10", lambda:
+       ((),
+        sparse_tensor.SparseTensor(indices=[[0]], values=[1], dense_shape=[1]),
+        ()), lambda: ((), sparse_tensor.SparseTensor, ())),
+      ("TestCase_11", lambda: ((), constant_op.constant([1]), ()), lambda:
+       ((), tensor.Tensor, ())),
+  ]
+
+  def reduce_fn(x, y):
+    name, classes_fn, expected_fn = y
+    return x + combinations.combine(
+        classes_fn=combinations.NamedObject("classes_fn.{}".format(name),
+                                            classes_fn),
+        expected_fn=combinations.NamedObject("expected_fn.{}".format(name),
+                                             expected_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _test_serialize_deserialize_combinations():
+  cases = [("TestCase_0", lambda: ()),
+           ("TestCase_1", lambda: sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1])),
+           ("TestCase_2", lambda: sparse_tensor.SparseTensor(
+               indices=[[3, 4]], values=[-1], dense_shape=[4, 5])),
+           ("TestCase_3", lambda: sparse_tensor.SparseTensor(
+               indices=[[0, 0], [3, 4]], values=[1, -1], dense_shape=[4, 5])),
+           ("TestCase_4", lambda: (sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1]))),
+           ("TestCase_5", lambda: (sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1]), ())),
+           ("TestCase_6", lambda:
+            ((),
+             sparse_tensor.SparseTensor(
+                 indices=[[0, 0]], values=[1], dense_shape=[1, 1])))]
+
+  def reduce_fn(x, y):
+    name, input_fn = y
+    return x + combinations.combine(
+        input_fn=combinations.NamedObject("input_fn.{}".format(name), input_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+def _test_serialize_many_deserialize_combinations():
+  cases = [("TestCase_0", lambda: ()),
+           ("TestCase_1", lambda: sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1])),
+           ("TestCase_2", lambda: sparse_tensor.SparseTensor(
+               indices=[[3, 4]], values=[-1], dense_shape=[4, 5])),
+           ("TestCase_3", lambda: sparse_tensor.SparseTensor(
+               indices=[[0, 0], [3, 4]], values=[1, -1], dense_shape=[4, 5])),
+           ("TestCase_4", lambda: (sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1]))),
+           ("TestCase_5", lambda: (sparse_tensor.SparseTensor(
+               indices=[[0, 0]], values=[1], dense_shape=[1, 1]), ())),
+           ("TestCase_6", lambda:
+            ((),
+             sparse_tensor.SparseTensor(
+                 indices=[[0, 0]], values=[1], dense_shape=[1, 1])))]
+
+  def reduce_fn(x, y):
+    name, input_fn = y
+    return x + combinations.combine(
+        input_fn=combinations.NamedObject("input_fn.{}".format(name), input_fn))
+
+  return functools.reduce(reduce_fn, cases, [])
+
+
+class SparseTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_any_sparse_combinations()))
+  def testAnySparse(self, classes_fn, expected):
+    classes = classes_fn()
+    self.assertEqual(sparse.any_sparse(classes), expected)
+
+  def assertShapesEqual(self, a, b):
+    for a, b in zip(nest.flatten(a), nest.flatten(b)):
+      self.assertEqual(a.ndims, b.ndims)
+      if a.ndims is None:
+        continue
+      for c, d in zip(a.as_list(), b.as_list()):
+        self.assertEqual(c, d)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_as_dense_shapes_combinations()))
+  def testAsDenseShapes(self, types_fn, classes_fn, expected_fn):
+    types = types_fn()
+    classes = classes_fn()
+    expected = expected_fn()
+    self.assertShapesEqual(sparse.as_dense_shapes(types, classes), expected)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_as_dense_types_combinations()))
+  def testAsDenseTypes(self, types_fn, classes_fn, expected_fn):
+    types = types_fn()
+    classes = classes_fn()
+    expected = expected_fn()
+    self.assertEqual(sparse.as_dense_types(types, classes), expected)
+
+  @combinations.generate(
+      combinations.times(test_base.default_test_combinations(),
+                         _test_get_classes_combinations()))
+  def testGetClasses(self, classes_fn, expected_fn):
+    classes = classes_fn()
+    expected = expected_fn()
+    self.assertEqual(sparse.get_classes(classes), expected)
+
+  def assertSparseValuesEqual(self, a, b):
+    if not isinstance(a, sparse_tensor.SparseTensor):
+      self.assertFalse(isinstance(b, sparse_tensor.SparseTensor))
+      self.assertEqual(a, b)
+      return
+    self.assertTrue(isinstance(b, sparse_tensor.SparseTensor))
+    with self.cached_session():
+      self.assertAllEqual(a.eval().indices, self.evaluate(b).indices)
+      self.assertAllEqual(a.eval().values, self.evaluate(b).values)
+      self.assertAllEqual(a.eval().dense_shape, self.evaluate(b).dense_shape)
+
+  @combinations.generate(
+      combinations.times(test_base.graph_only_combinations(),
+                         _test_serialize_deserialize_combinations()))
+  def testSerializeDeserialize(self, input_fn):
+    test_case = input_fn()
+    classes = sparse.get_classes(test_case)
+    shapes = nest.map_structure(lambda _: tensor_shape.TensorShape(None),
+                                classes)
+    types = nest.map_structure(lambda _: dtypes.int32, classes)
+    actual = sparse.deserialize_sparse_tensors(
+        sparse.serialize_sparse_tensors(test_case), types, shapes,
+        sparse.get_classes(test_case))
+    nest.assert_same_structure(test_case, actual)
+    for a, e in zip(nest.flatten(actual), nest.flatten(test_case)):
+      self.assertSparseValuesEqual(a, e)
+
+  @combinations.generate(
+      combinations.times(test_base.graph_only_combinations(),
+                         _test_serialize_many_deserialize_combinations()))
+  def testSerializeManyDeserialize(self, input_fn):
+    test_case = input_fn()
+    classes = sparse.get_classes(test_case)
+    shapes = nest.map_structure(lambda _: tensor_shape.TensorShape(None),
+                                classes)
+    types = nest.map_structure(lambda _: dtypes.int32, classes)
+    actual = sparse.deserialize_sparse_tensors(
+        sparse.serialize_many_sparse_tensors(test_case), types, shapes,
+        sparse.get_classes(test_case))
+    nest.assert_same_structure(test_case, actual)
+    for a, e in zip(nest.flatten(actual), nest.flatten(test_case)):
+      self.assertSparseValuesEqual(a, e)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/util/structure.py

@@ -0,0 +1,521 @@
+# 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.
+# ==============================================================================
+"""Utilities for describing the structure of a `tf.data` type."""
+import collections
+import functools
+import itertools
+
+import wrapt
+
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import none_tensor
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.types import internal
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.compat import collections_abc
+from tensorflow.python.util.nest_util import CustomNestProtocol
+from tensorflow.python.util.tf_export import tf_export
+
+
+# pylint: disable=invalid-name
+@tf_export(v1=["data.experimental.TensorStructure"])
+@deprecation.deprecated(None, "Use `tf.TensorSpec` instead.")
+def _TensorStructure(dtype, shape):
+  return tensor_lib.TensorSpec(shape, dtype)
+
+
+@tf_export(v1=["data.experimental.SparseTensorStructure"])
+@deprecation.deprecated(None, "Use `tf.SparseTensorSpec` instead.")
+def _SparseTensorStructure(dtype, shape):
+  return sparse_tensor.SparseTensorSpec(shape, dtype)
+
+
+@tf_export(v1=["data.experimental.TensorArrayStructure"])
+@deprecation.deprecated(None, "Use `tf.TensorArraySpec` instead.")
+def _TensorArrayStructure(dtype, element_shape, dynamic_size, infer_shape):
+  return tensor_array_ops.TensorArraySpec(element_shape, dtype,
+                                          dynamic_size, infer_shape)
+
+
+@tf_export(v1=["data.experimental.RaggedTensorStructure"])
+@deprecation.deprecated(None, "Use `tf.RaggedTensorSpec` instead.")
+def _RaggedTensorStructure(dtype, shape, ragged_rank):
+  return ragged_tensor.RaggedTensorSpec(shape, dtype, ragged_rank)
+# pylint: enable=invalid-name
+
+
+# TODO(jsimsa): Remove the special-case for `TensorArray` pass-through once
+# it is a subclass of `CompositeTensor`.
+def normalize_element(element, element_signature=None):
+  """Normalizes a nested structure of element components.
+
+  * Components matching `SparseTensorSpec` are converted to `SparseTensor`.
+  * Components matching `RaggedTensorSpec` are converted to `RaggedTensor`.
+  * Components matching `VariableSpec` are converted to `Tensor`.
+  * Components matching `DatasetSpec` or `TensorArraySpec` are passed through.
+  * `CompositeTensor` components are passed through.
+  * All other components are converted to `Tensor`.
+
+  Args:
+    element: A nested structure of individual components.
+    element_signature: (Optional.) A nested structure of `tf.DType` objects
+      corresponding to each component of `element`. If specified, it will be
+      used to set the exact type of output tensor when converting input
+      components which are not tensors themselves (e.g. numpy arrays, native
+      python types, etc.)
+
+  Returns:
+    A nested structure of `Tensor`, `Variable`, `Dataset`, `SparseTensor`,
+    `RaggedTensor`, or `TensorArray` objects.
+  """
+  normalized_components = []
+  if element_signature is None:
+    components = nest.flatten(element)
+    flattened_signature = [None] * len(components)
+    pack_as = element
+  else:
+    flattened_signature = nest.flatten(element_signature)
+    components = nest.flatten_up_to(element_signature, element)
+    pack_as = element_signature
+  with ops.name_scope("normalize_element"):
+    for i, (t, spec) in enumerate(zip(components, flattened_signature)):
+      try:
+        if spec is None:
+          spec = type_spec_from_value(t, use_fallback=False)
+      except TypeError:
+        # TypeError indicates it was not possible to compute a `TypeSpec` for
+        # the value. As a fallback try converting the value to a tensor.
+        normalized_components.append(
+            ops.convert_to_tensor(t, name="component_%d" % i))
+      else:
+        # To avoid a circular dependency between dataset_ops and structure,
+        # we check the class name instead of using `isinstance`.
+        if spec.__class__.__name__ == "DatasetSpec":
+          normalized_components.append(t)
+        elif isinstance(spec, sparse_tensor.SparseTensorSpec):
+          normalized_components.append(sparse_tensor.SparseTensor.from_value(t))
+        elif isinstance(spec, ragged_tensor.RaggedTensorSpec):
+          normalized_components.append(
+              ragged_tensor.convert_to_tensor_or_ragged_tensor(
+                  t, name="component_%d" % i))
+        elif isinstance(spec, (tensor_array_ops.TensorArraySpec)):
+          normalized_components.append(t)
+        elif isinstance(spec, none_tensor.NoneTensorSpec):
+          normalized_components.append(none_tensor.NoneTensor())
+        elif isinstance(spec, resource_variable_ops.VariableSpec):
+          normalized_components.append(
+              ops.convert_to_tensor(t, name=f"component_{i}", dtype=spec.dtype))
+        elif isinstance(t, composite_tensor.CompositeTensor):
+          normalized_components.append(t)
+        else:
+          dtype = getattr(spec, "dtype", None)
+          normalized_components.append(
+              ops.convert_to_tensor(t, name="component_%d" % i, dtype=dtype))
+  return nest.pack_sequence_as(pack_as, normalized_components)
+
+
+def convert_legacy_structure(output_types, output_shapes, output_classes):
+  """Returns a `Structure` that represents the given legacy structure.
+
+  This method provides a way to convert from the existing `Dataset` and
+  `Iterator` structure-related properties to a `Structure` object. A "legacy"
+  structure is represented by the `tf.data.Dataset.output_types`,
+  `tf.data.Dataset.output_shapes`, and `tf.data.Dataset.output_classes`
+  properties.
+
+  TODO(b/110122868): Remove this function once `Structure` is used throughout
+  `tf.data`.
+
+  Args:
+    output_types: A nested structure of `tf.DType` objects corresponding to
+      each component of a structured value.
+    output_shapes: A nested structure of `tf.TensorShape` objects
+      corresponding to each component a structured value.
+    output_classes: A nested structure of Python `type` objects corresponding
+      to each component of a structured value.
+
+  Returns:
+    A `Structure`.
+
+  Raises:
+    TypeError: If a structure cannot be built from the arguments, because one of
+      the component classes in `output_classes` is not supported.
+  """
+  flat_types = nest.flatten(output_types)
+  flat_shapes = nest.flatten(output_shapes)
+  flat_classes = nest.flatten(output_classes)
+  flat_ret = []
+  for flat_type, flat_shape, flat_class in zip(flat_types, flat_shapes,
+                                               flat_classes):
+    if isinstance(flat_class, type_spec.TypeSpec):
+      flat_ret.append(flat_class)
+    elif issubclass(flat_class, sparse_tensor.SparseTensor):
+      flat_ret.append(sparse_tensor.SparseTensorSpec(flat_shape, flat_type))
+    elif issubclass(flat_class, tensor_lib.Tensor):
+      flat_ret.append(tensor_lib.TensorSpec(flat_shape, flat_type))
+    elif issubclass(flat_class, tensor_array_ops.TensorArray):
+      # We sneaked the dynamic_size and infer_shape into the legacy shape.
+      flat_ret.append(
+          tensor_array_ops.TensorArraySpec(
+              flat_shape[2:], flat_type,
+              dynamic_size=tensor_shape.dimension_value(flat_shape[0]),
+              infer_shape=tensor_shape.dimension_value(flat_shape[1])))
+    else:
+      # NOTE(mrry): Since legacy structures produced by iterators only
+      # comprise Tensors, SparseTensors, and nests, we do not need to
+      # support all structure types here.
+      raise TypeError(
+          "Could not build a structure for output class {}. Make sure any "
+          "component class in `output_classes` inherits from one of the "
+          "following classes: `tf.TypeSpec`, `tf.sparse.SparseTensor`, "
+          "`tf.Tensor`, `tf.TensorArray`.".format(flat_class.__name__))
+
+  return nest.pack_sequence_as(output_classes, flat_ret)
+
+
+def _from_tensor_list_helper(decode_fn, element_spec, tensor_list):
+  """Returns an element constructed from the given spec and tensor list.
+
+  Args:
+    decode_fn: Method that constructs an element component from the element spec
+      component and a tensor list.
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    tensor_list: A list of tensors to use for constructing the value.
+
+  Returns:
+    An element constructed from the given spec and tensor list.
+
+  Raises:
+    ValueError: If the number of tensors needed to construct an element for
+      the given spec does not match the given number of tensors.
+  """
+
+  # pylint: disable=protected-access
+
+  flat_specs = nest.flatten(element_spec)
+  flat_spec_lengths = [len(spec._flat_tensor_specs) for spec in flat_specs]
+  if sum(flat_spec_lengths) != len(tensor_list):
+    raise ValueError("Expected {} tensors but got {}.".format(
+        sum(flat_spec_lengths), len(tensor_list)))
+
+  i = 0
+  flat_ret = []
+  for (component_spec, num_flat_values) in zip(flat_specs, flat_spec_lengths):
+    value = tensor_list[i:i + num_flat_values]
+    flat_ret.append(decode_fn(component_spec, value))
+    i += num_flat_values
+  return nest.pack_sequence_as(element_spec, flat_ret)
+
+
+def from_compatible_tensor_list(element_spec, tensor_list):
+  """Returns an element constructed from the given spec and tensor list.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    tensor_list: A list of tensors to use for constructing the value.
+
+  Returns:
+    An element constructed from the given spec and tensor list.
+
+  Raises:
+    ValueError: If the number of tensors needed to construct an element for
+      the given spec does not match the given number of tensors.
+  """
+
+  # pylint: disable=protected-access
+  # pylint: disable=g-long-lambda
+  return _from_tensor_list_helper(
+      lambda spec, value: spec._from_compatible_tensor_list(value),
+      element_spec, tensor_list)
+
+
+def from_tensor_list(element_spec, tensor_list):
+  """Returns an element constructed from the given spec and tensor list.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    tensor_list: A list of tensors to use for constructing the value.
+
+  Returns:
+    An element constructed from the given spec and tensor list.
+
+  Raises:
+    ValueError: If the number of tensors needed to construct an element for
+      the given spec does not match the given number of tensors or the given
+      spec is not compatible with the tensor list.
+  """
+
+  # pylint: disable=protected-access
+  # pylint: disable=g-long-lambda
+  return _from_tensor_list_helper(
+      lambda spec, value: spec._from_tensor_list(value), element_spec,
+      tensor_list)
+
+
+def get_flat_tensor_specs(element_spec):
+  """Returns a list `tf.TypeSpec`s for the element tensor representation.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+
+  Returns:
+    A list `tf.TypeSpec`s for the element tensor representation.
+  """
+
+  # pylint: disable=protected-access
+  return list(
+      itertools.chain.from_iterable(
+          spec._flat_tensor_specs for spec in nest.flatten(element_spec)))
+
+
+def get_flat_tensor_shapes(element_spec):
+  """Returns a list `tf.TensorShapes`s for the element tensor representation.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+
+  Returns:
+    A list `tf.TensorShapes`s for the element tensor representation.
+  """
+  return [spec.shape for spec in get_flat_tensor_specs(element_spec)]
+
+
+def get_flat_tensor_types(element_spec):
+  """Returns a list `tf.DType`s for the element tensor representation.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+
+  Returns:
+    A list `tf.DType`s for the element tensor representation.
+  """
+  return [spec.dtype for spec in get_flat_tensor_specs(element_spec)]
+
+
+def _to_tensor_list_helper(encode_fn, element_spec, element):
+  """Returns a tensor list representation of the element.
+
+  Args:
+    encode_fn: Method that constructs a tensor list representation from the
+      given element spec and element.
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    element: The element to convert to tensor list representation.
+
+  Returns:
+    A tensor list representation of `element`.
+
+  Raises:
+    ValueError: If `element_spec` and `element` do not have the same number of
+      elements or if the two structures are not nested in the same way.
+    TypeError: If `element_spec` and `element` differ in the type of sequence
+      in any of their substructures.
+  """
+
+  nest.assert_same_structure(element_spec, element)
+
+  def reduce_fn(state, value):
+    spec, component = value
+    if isinstance(spec, internal.TensorSpec):
+      try:
+        component = ops.convert_to_tensor(component, spec.dtype)
+      except (TypeError, ValueError):
+        raise ValueError(
+            f"Value {component} is not convertible to a tensor with "
+            f"dtype {spec.dtype} and shape {spec.shape}."
+        )
+      if not component.shape.is_compatible_with(spec.shape):
+        raise ValueError(
+            f"Value {component} is not convertible to a tensor with "
+            f"dtype {spec.dtype} and shape {spec.shape}."
+        )
+    return encode_fn(state, spec, component)
+
+  return functools.reduce(
+      reduce_fn, zip(nest.flatten(element_spec), nest.flatten(element)), [])
+
+
+def to_batched_tensor_list(element_spec, element):
+  """Returns a tensor list representation of the element.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    element: The element to convert to tensor list representation.
+
+  Returns:
+    A tensor list representation of `element`.
+
+  Raises:
+    ValueError: If `element_spec` and `element` do not have the same number of
+      elements or if the two structures are not nested in the same way or the
+      rank of any of the tensors in the tensor list representation is 0.
+    TypeError: If `element_spec` and `element` differ in the type of sequence
+      in any of their substructures.
+  """
+
+  # pylint: disable=protected-access
+  # pylint: disable=g-long-lambda
+  return _to_tensor_list_helper(
+      lambda state, spec, component: state + spec._to_batched_tensor_list(
+          component), element_spec, element)
+
+
+def to_tensor_list(element_spec, element):
+  """Returns a tensor list representation of the element.
+
+  Args:
+    element_spec: A nested structure of `tf.TypeSpec` objects representing to
+      element type specification.
+    element: The element to convert to tensor list representation.
+
+  Returns:
+    A tensor list representation of `element`.
+
+  Raises:
+    ValueError: If `element_spec` and `element` do not have the same number of
+      elements or if the two structures are not nested in the same way.
+    TypeError: If `element_spec` and `element` differ in the type of sequence
+      in any of their substructures.
+  """
+
+  # pylint: disable=protected-access
+  # pylint: disable=g-long-lambda
+  return _to_tensor_list_helper(
+      lambda state, spec, component: state + spec._to_tensor_list(component),
+      element_spec, element)
+
+
+def are_compatible(spec1, spec2):
+  """Indicates whether two type specifications are compatible.
+
+  Two type specifications are compatible if they have the same nested structure
+  and the their individual components are pair-wise compatible.
+
+  Args:
+    spec1: A `tf.TypeSpec` object to compare.
+    spec2: A `tf.TypeSpec` object to compare.
+
+  Returns:
+    `True` if the two type specifications are compatible and `False` otherwise.
+  """
+
+  try:
+    nest.assert_same_structure(spec1, spec2)
+  except TypeError:
+    return False
+  except ValueError:
+    return False
+
+  for s1, s2 in zip(nest.flatten(spec1), nest.flatten(spec2)):
+    if not s1.is_compatible_with(s2) or not s2.is_compatible_with(s1):
+      return False
+  return True
+
+
+def type_spec_from_value(element, use_fallback=True):
+  """Creates a type specification for the given value.
+
+  Args:
+    element: The element to create the type specification for.
+    use_fallback: Whether to fall back to converting the element to a tensor
+      in order to compute its `TypeSpec`.
+
+  Returns:
+    A nested structure of `TypeSpec`s that represents the type specification
+    of `element`.
+
+  Raises:
+    TypeError: If a `TypeSpec` cannot be built for `element`, because its type
+      is not supported.
+  """
+  spec = type_spec._type_spec_from_value(element)  # pylint: disable=protected-access
+  if spec is not None:
+    return spec
+
+  if isinstance(element, collections_abc.Mapping):
+    # We create a shallow copy in an attempt to preserve the key order.
+    #
+    # Note that we do not guarantee that the key order is preserved, which is
+    # a limitation inherited from `copy()`. As a consequence, callers of
+    # `type_spec_from_value` should not assume that the key order of a `dict`
+    # in the returned nested structure matches the key order of the
+    # corresponding `dict` in the input value.
+    if isinstance(element, collections.defaultdict):
+      ctor = lambda items: type(element)(element.default_factory, items)
+    else:
+      ctor = type(element)
+    return ctor([(k, type_spec_from_value(v)) for k, v in element.items()])
+
+  if isinstance(element, tuple):
+    if hasattr(element, "_fields") and isinstance(
+        element._fields, collections_abc.Sequence) and all(
+            isinstance(f, str) for f in element._fields):
+      if isinstance(element, wrapt.ObjectProxy):
+        element_type = type(element.__wrapped__)
+      else:
+        element_type = type(element)
+      # `element` is a namedtuple
+      return element_type(*[type_spec_from_value(v) for v in element])
+    # `element` is not a namedtuple
+    return tuple([type_spec_from_value(v) for v in element])
+
+  if hasattr(element.__class__, "__attrs_attrs__"):
+    # `element` is an `attr.s` decorated class
+    attrs = getattr(element.__class__, "__attrs_attrs__")
+    return type(element)(*[
+        type_spec_from_value(getattr(element, a.name)) for a in attrs
+    ])
+
+  if isinstance(element, CustomNestProtocol):
+    # pylint: disable=protected-access
+    metadata, children = element.__tf_flatten__()
+    return element.__tf_unflatten__(metadata, type_spec_from_value(children))
+    # pylint: enable=protected-access
+
+  if use_fallback:
+    # As a fallback try converting the element to a tensor.
+    try:
+      tensor = ops.convert_to_tensor(element)
+      spec = type_spec_from_value(tensor)
+      if spec is not None:
+        return spec
+    except (ValueError, TypeError) as e:
+      logging.vlog(
+          3, "Failed to convert %r to tensor: %s" % (type(element).__name__, e))
+
+  raise TypeError("Could not build a `TypeSpec` for {} with type {}".format(
+      element,
+      type(element).__name__))
+
+
+# TODO(b/149584798): remove legacy forwarding references
+NoneTensor = none_tensor.NoneTensor
+NoneTensorSpec = none_tensor.NoneTensorSpec

tensorflow/python/data/util/traverse.py

@@ -0,0 +1,81 @@
+# 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.
+# ==============================================================================
+"""Helpers to traverse the Dataset dependency structure."""
+import queue
+
+from tensorflow.python.framework import dtypes
+
+
+OP_TYPES_ALLOWLIST = ["DummyIterationCounter"]
+# We allowlist all ops that produce variant tensors as output. This is a bit
+# of overkill but the other dataset _inputs() traversal strategies can't
+# cover the case of function inputs that capture dataset variants.
+TENSOR_TYPES_ALLOWLIST = [dtypes.variant]
+
+
+def _traverse(dataset, op_filter_fn):
+  """Traverse a dataset graph, returning nodes matching `op_filter_fn`."""
+  result = []
+  bfs_q = queue.Queue()
+  bfs_q.put(dataset._variant_tensor.op)  # pylint: disable=protected-access
+  visited = []
+  while not bfs_q.empty():
+    op = bfs_q.get()
+    visited.append(op)
+    if op_filter_fn(op):
+      result.append(op)
+    for i in op.inputs:
+      input_op = i.op
+      if input_op not in visited:
+        bfs_q.put(input_op)
+  return result
+
+
+def obtain_capture_by_value_ops(dataset):
+  """Given an input dataset, finds all allowlisted ops used for construction.
+
+  Allowlisted ops are stateful ops which are known to be safe to capture by
+  value.
+
+  Args:
+    dataset: Dataset to find allowlisted stateful ops for.
+
+  Returns:
+    A list of variant_tensor producing dataset ops used to construct this
+    dataset.
+  """
+
+  def capture_by_value(op):
+    return (op.outputs[0].dtype in TENSOR_TYPES_ALLOWLIST or
+            op.type in OP_TYPES_ALLOWLIST)
+
+  return _traverse(dataset, capture_by_value)
+
+
+def obtain_all_variant_tensor_ops(dataset):
+  """Given an input dataset, finds all dataset ops used for construction.
+
+  A series of transformations would have created this dataset with each
+  transformation including zero or more Dataset ops, each producing a dataset
+  variant tensor. This method outputs all of them.
+
+  Args:
+    dataset: Dataset to find variant tensors for.
+
+  Returns:
+    A list of variant_tensor producing dataset ops used to construct this
+    dataset.
+  """
+  return _traverse(dataset, lambda op: op.outputs[0].dtype == dtypes.variant)

tensorflow/python/data/util/traverse_test.py

@@ -0,0 +1,123 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for utilities for traversing the dataset construction graph."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import compat
+from tensorflow.python.data.experimental.ops import data_service_ops
+from tensorflow.python.data.kernel_tests import test_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.util import traverse
+from tensorflow.python.framework import combinations
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class _TestDataset(dataset_ops.UnaryUnchangedStructureDataset):
+
+  def __init__(self, input_dataset):
+    self._input_dataset = input_dataset
+    temp_variant_tensor = gen_dataset_ops.prefetch_dataset(
+        input_dataset._variant_tensor,
+        buffer_size=1,
+        **self._flat_structure)
+    variant_tensor = gen_dataset_ops.model_dataset(
+        temp_variant_tensor, **self._flat_structure)
+    super(_TestDataset, self).__init__(input_dataset, variant_tensor)
+
+
+class TraverseTest(test_base.DatasetTestBase, parameterized.TestCase):
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testOnlySource(self):
+    ds = dataset_ops.Dataset.range(10)
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
+    self.assertAllEqual(["RangeDataset"], [x.name for x in variant_tensor_ops])
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testSimplePipeline(self):
+    ds = dataset_ops.Dataset.range(10).map(math_ops.square)
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
+    self.assertSetEqual(
+        set(["MapDataset", "RangeDataset"]),
+        set(x.name for x in variant_tensor_ops))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testConcat(self):
+    ds1 = dataset_ops.Dataset.range(10)
+    ds2 = dataset_ops.Dataset.range(10)
+    ds = ds1.concatenate(ds2)
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
+    self.assertSetEqual(
+        set(["ConcatenateDataset", "RangeDataset", "RangeDataset_1"]),
+        set(x.name for x in variant_tensor_ops))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testZip(self):
+    ds1 = dataset_ops.Dataset.range(10)
+    ds2 = dataset_ops.Dataset.range(10)
+    ds = dataset_ops.Dataset.zip((ds1, ds2))
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
+    self.assertSetEqual(
+        set(["ZipDataset", "RangeDataset", "RangeDataset_1"]),
+        set(x.name for x in variant_tensor_ops))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testMultipleVariantTensors(self):
+    ds = dataset_ops.Dataset.range(10)
+    ds = _TestDataset(ds)
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds)
+    self.assertSetEqual(
+        set(["RangeDataset", "ModelDataset", "PrefetchDataset"]),
+        set(x.name for x in variant_tensor_ops))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testFlatMap(self):
+    ds1 = dataset_ops.Dataset.range(10).repeat(10)
+
+    def map_fn(ds):
+
+      def _map(x):
+        return ds.batch(x)
+
+      return _map
+
+    ds2 = dataset_ops.Dataset.range(20).prefetch(1)
+    ds2 = ds2.flat_map(map_fn(ds1))
+    variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(ds2)
+    self.assertSetEqual(
+        set([
+            "FlatMapDataset", "PrefetchDataset", "RepeatDataset",
+            "RangeDataset", "RangeDataset_1"
+        ]), set(x.name for x in variant_tensor_ops))
+
+  @combinations.generate(test_base.graph_only_combinations())
+  def testTfDataService(self):
+    ds = dataset_ops.Dataset.range(10)
+    ds = ds.apply(
+        data_service_ops.distribute("parallel_epochs", "grpc://foo:0"))
+    ops = traverse.obtain_capture_by_value_ops(ds)
+    data_service_dataset_op = ("DataServiceDatasetV4"
+                               if compat.forward_compatible(2022, 8, 31) else
+                               "DataServiceDatasetV3")
+    self.assertContainsSubset(
+        ["RangeDataset", data_service_dataset_op, "DummyIterationCounter"],
+        set(x.name for x in ops))
+
+
+if __name__ == "__main__":
+  test.main()