# 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()

0f68e419f965bd0df0bf0ef7c1788f27c46e21b4
<mask token>


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
            )
    <mask token>

    @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)


<mask token>

<mask token>


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, [])


<mask token>


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, [])


<mask token>


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)


<mask token>

<mask token>


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()),
        ('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()),
        ('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, [])


<mask token>


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, [])


<mask token>


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)


<mask token>

<mask token>


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()),
        ('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()),
        ('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)


<mask token>

# 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()
