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