diff --git a/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py b/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py index 1bf07a98f28..1240b704119 100644 --- a/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/matching_files_test.py @@ -34,10 +34,12 @@ class MatchingFilesDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): def setUp(self): + super(MatchingFilesDatasetTest, self).setUp() self.tmp_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmp_dir, ignore_errors=True) + super(MatchingFilesDatasetTest, self).tearDown() def _touchTempFiles(self, filenames): for filename in filenames: diff --git a/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py index f6ab5a1cde2..aea4934260e 100644 --- a/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py +++ b/tensorflow/python/data/experimental/kernel_tests/serialization/dataset_serialization_test_base.py @@ -57,6 +57,7 @@ class DatasetSerializationTestBase(test.TestCase): def tearDown(self): self._delete_ckpt() + super(DatasetSerializationTestBase, self).tearDown() # TODO(b/72657739): Remove sparse_tensor argument, which is to test the # (deprecated) saveable `SparseTensorSliceDataset`, once the API diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD index 2b1bda4138a..db749da77f8 100644 --- a/tensorflow/python/data/kernel_tests/BUILD +++ b/tensorflow/python/data/kernel_tests/BUILD @@ -287,6 +287,7 @@ tf_py_test( size = "small", srcs = ["iterator_cluster_test.py"], additional_deps = [ + ":test_base", "//tensorflow/core:protos_all_py", "//tensorflow/python/data/ops:dataset_ops", "//tensorflow/python/data/ops:iterator_ops", @@ -400,6 +401,7 @@ tf_py_test( "//tensorflow/python:variable_scope", "//tensorflow/python/ops/ragged", ], + shard_count = 4, ) cuda_py_test( diff --git a/tensorflow/python/data/kernel_tests/as_numpy_iterator_test.py b/tensorflow/python/data/kernel_tests/as_numpy_iterator_test.py index 23a928a817c..b704906a3ae 100644 --- a/tensorflow/python/data/kernel_tests/as_numpy_iterator_test.py +++ b/tensorflow/python/data/kernel_tests/as_numpy_iterator_test.py @@ -61,30 +61,30 @@ class AsNumpyIteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(RuntimeError): ds.as_numpy_iterator() - def checkInvalidElement(self, element): + def _testInvalidElement(self, element): ds = dataset_ops.Dataset.from_tensors(element) with self.assertRaisesRegex(TypeError, '.*does not support datasets containing.*'): ds.as_numpy_iterator() @combinations.generate(test_base.eager_only_combinations()) - def testInvalidElements(self): - self.checkInvalidElement(sparse_tensor.SparseTensorValue([[0]], [0], [1])) + def testSparseElement(self): + self._testInvalidElement(sparse_tensor.SparseTensorValue([[0]], [0], [1])) @combinations.generate(test_base.eager_only_combinations()) def testRaggedElement(self): - self.checkInvalidElement( + self._testInvalidElement( ragged_tensor_value.RaggedTensorValue( np.array([0, 1, 2]), np.array([0, 1, 3], dtype=np.int64))) @combinations.generate(test_base.eager_only_combinations()) def testDatasetElement(self): - self.checkInvalidElement(dataset_ops.Dataset.range(3)) + self._testInvalidElement(dataset_ops.Dataset.range(3)) @combinations.generate(test_base.eager_only_combinations()) def testNestedNonTensorElement(self): tuple_elem = (constant_op.constant([1, 2, 3]), dataset_ops.Dataset.range(3)) - self.checkInvalidElement(tuple_elem) + self._testInvalidElement(tuple_elem) if __name__ == '__main__': diff --git a/tensorflow/python/data/kernel_tests/cache_test.py b/tensorflow/python/data/kernel_tests/cache_test.py index 6d3dc04a3e0..1a923645a04 100644 --- a/tensorflow/python/data/kernel_tests/cache_test.py +++ b/tensorflow/python/data/kernel_tests/cache_test.py @@ -45,9 +45,9 @@ class FileCacheTest(test_base.DatasetTestBase, parameterized.TestCase): self.cache_prefix = path.join(self.tmp_dir, "cache") def tearDown(self): - super(FileCacheTest, self).tearDown() if self.tmp_dir: shutil.rmtree(self.tmp_dir, ignore_errors=True) + super(FileCacheTest, self).tearDown() @combinations.generate(test_base.default_test_combinations()) def testCacheDatasetPassthrough(self): diff --git a/tensorflow/python/data/kernel_tests/checkpoint_test.py b/tensorflow/python/data/kernel_tests/checkpoint_test.py index 738d09b97fe..4441d5642d5 100644 --- a/tensorflow/python/data/kernel_tests/checkpoint_test.py +++ b/tensorflow/python/data/kernel_tests/checkpoint_test.py @@ -42,11 +42,11 @@ from tensorflow.python.training.tracking import util as trackable_utils class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): def tearDown(self): - super(CheckpointTest, self).tearDown() prefix = self._iterator_checkpoint_prefix() pattern = prefix + "*" files = gfile.Glob(pattern) map(gfile.Remove, files) + super(CheckpointTest, self).tearDown() def _iterator_checkpoint_prefix(self): return os.path.join(self.get_temp_dir(), "iterator") @@ -66,8 +66,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): iterator_state_variant) return restore_op - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="graph")) + @combinations.generate(test_base.graph_only_combinations()) def testSaveRestore(self): def _build_graph(start, stop): @@ -118,8 +117,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="graph")) + @combinations.generate(test_base.graph_only_combinations()) def testInitThenRestore(self): # Note: Calling init_op before restore_op is redundant. This test just makes # sure we do not fail if restore is called on an already initialized @@ -157,8 +155,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="graph")) + @combinations.generate(test_base.graph_only_combinations()) def testMultipleSaves(self): def _build_graph(start, stop): @@ -204,8 +201,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="graph")) + @combinations.generate(test_base.graph_only_combinations()) def testSaveRestoreWithRepeat(self): def _build_graph(start, stop, num_epochs): @@ -253,8 +249,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="graph")) + @combinations.generate(test_base.graph_only_combinations()) def testSaveRestoreExhaustedIterator(self): def _build_graph(start, stop, num_epochs): @@ -295,8 +290,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testSaveRestoreOneShotIterator(self): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") @@ -319,8 +313,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): get_next() - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testSaveRestoreMultipleIterator(self): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") @@ -353,8 +346,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertAllEqual([1, 4], get_next_2()) self.assertAllEqual(3, get_next_3()) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testRestoreExhaustedIterator(self): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") @@ -373,8 +365,7 @@ class CheckpointTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): get_next() - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testRestoreInReconstructedIteratorInitializable(self): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") diff --git a/tensorflow/python/data/kernel_tests/dataset_test.py b/tensorflow/python/data/kernel_tests/dataset_test.py index b35c4ff1b29..df151a85db0 100644 --- a/tensorflow/python/data/kernel_tests/dataset_test.py +++ b/tensorflow/python/data/kernel_tests/dataset_test.py @@ -43,7 +43,6 @@ from tensorflow.python.framework import tensor_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import random_ops from tensorflow.python.platform import test -from tensorflow.python.platform import tf_logging class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @@ -89,13 +88,13 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): variant, original_dataset.element_spec) self.assertDatasetProduces(revived_dataset, list(original_dataset)) - def checkNumInputs(self, dataset, num_inputs): + def _testNumInputs(self, dataset, num_inputs): self.assertLen(dataset._inputs(), num_inputs) @combinations.generate(test_base.default_test_combinations()) def testFixedLengthRecordInputs(self): dataset = readers.FixedLengthRecordDataset("", 42) - self.checkNumInputs(dataset, 0) + self._testNumInputs(dataset, 0) @combinations.generate(test_base.default_test_combinations()) def testFromGeneratorInputs(self): @@ -103,27 +102,27 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): yield 42 dataset = dataset_ops.Dataset.from_generator(gen, dtypes.int32) - self.checkNumInputs(dataset, 1) + self._testNumInputs(dataset, 1) @combinations.generate(test_base.default_test_combinations()) def testFromTensorsInputs(self): dataset = dataset_ops.Dataset.from_tensors([42]) - self.checkNumInputs(dataset, 0) + self._testNumInputs(dataset, 0) @combinations.generate(test_base.default_test_combinations()) def testRangeInputs(self): dataset = dataset_ops.Dataset.range(10) - self.checkNumInputs(dataset, 0) + self._testNumInputs(dataset, 0) @combinations.generate(test_base.default_test_combinations()) def testTextLineInputs(self): dataset = readers.TextLineDataset("") - self.checkNumInputs(dataset, 0) + self._testNumInputs(dataset, 0) @combinations.generate(test_base.default_test_combinations()) def testTFRecordInputs(self): dataset = readers.TFRecordDataset("") - self.checkNumInputs(dataset, 1) + self._testNumInputs(dataset, 1) @combinations.generate( combinations.combine(tf_api_version=1, mode=["eager", "graph"])) @@ -135,58 +134,58 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): dense_shape=np.array([3, 1]))) self.assertEmpty(dataset_fn._inputs()) - def checkUnaryInputs(self, dataset_fn): + def _testUnaryInputs(self, dataset_fn): input_dataset = dataset_ops.Dataset.range(0) self.assertEqual([input_dataset], dataset_fn(input_dataset)._inputs()) @combinations.generate(test_base.default_test_combinations()) def testBatchInputs(self): - self.checkUnaryInputs(lambda x: x.batch(10)) + self._testUnaryInputs(lambda x: x.batch(10)) @combinations.generate(test_base.default_test_combinations()) def testCacheInputs(self): - self.checkUnaryInputs(lambda x: x.cache()) + self._testUnaryInputs(lambda x: x.cache()) @combinations.generate(test_base.default_test_combinations()) def testFilterInputs(self): - self.checkUnaryInputs(lambda x: x.filter(lambda x: True)) + self._testUnaryInputs(lambda x: x.filter(lambda x: True)) @combinations.generate(test_base.default_test_combinations()) def testFlatMapInputs(self): - self.checkUnaryInputs( + self._testUnaryInputs( lambda x: x.flat_map(lambda x: dataset_ops.Dataset.range(0))) @combinations.generate(test_base.default_test_combinations()) def testMapInputs(self): - self.checkUnaryInputs(lambda x: x.map(lambda x: x)) + self._testUnaryInputs(lambda x: x.map(lambda x: x)) @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchInputs(self): - self.checkUnaryInputs(lambda x: x.padded_batch(10, [])) + self._testUnaryInputs(lambda x: x.padded_batch(10, [])) @combinations.generate(test_base.default_test_combinations()) def testParallelMapInputs(self): - self.checkUnaryInputs(lambda x: x.map(lambda x: x, num_parallel_calls=2)) + self._testUnaryInputs(lambda x: x.map(lambda x: x, num_parallel_calls=2)) @combinations.generate(test_base.default_test_combinations()) def testRepeatInputs(self): - self.checkUnaryInputs(lambda x: x.repeat()) + self._testUnaryInputs(lambda x: x.repeat()) @combinations.generate(test_base.default_test_combinations()) def testShuffleInputs(self): - self.checkUnaryInputs(lambda x: x.shuffle(10)) + self._testUnaryInputs(lambda x: x.shuffle(10)) @combinations.generate(test_base.default_test_combinations()) def testSkipInputs(self): - self.checkUnaryInputs(lambda x: x.skip(1)) + self._testUnaryInputs(lambda x: x.skip(1)) @combinations.generate(test_base.default_test_combinations()) def testTakeInputs(self): - self.checkUnaryInputs(lambda x: x.take(1)) + self._testUnaryInputs(lambda x: x.take(1)) @combinations.generate(test_base.default_test_combinations()) def testWindowInputs(self): - self.checkUnaryInputs(lambda x: x.window(10)) + self._testUnaryInputs(lambda x: x.window(10)) @combinations.generate(test_base.default_test_combinations()) def testUnaryTransformationInputsApply(self): @@ -195,7 +194,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual([input_dataset], dataset._inputs()) - def checkInputsWithInterleaveFn(self, dataset_fn, interleave_parallelism): + def _testInputsWithInterleaveFn(self, dataset_fn, interleave_parallelism): input_dataset = dataset_ops.Dataset.range(0) dataset = input_dataset.interleave( lambda x: dataset_ops.Dataset.range(0), @@ -205,11 +204,11 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) def testParallelInterleaveInputs(self): - self.checkInputsWithInterleaveFn(lambda: dataset_ops.range(0), 2) + self._testInputsWithInterleaveFn(lambda: dataset_ops.range(0), 2) @combinations.generate(test_base.default_test_combinations()) def testInterleaveInputs(self): - self.checkInputsWithInterleaveFn(lambda: dataset_ops.range(0), None) + self._testInputsWithInterleaveFn(lambda: dataset_ops.range(0), None) @combinations.generate(test_base.default_test_combinations()) def testNoWarnings(self): @@ -218,16 +217,16 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): lambda x: dataset_ops.Dataset.range(0), cycle_length=2) self.assertEmpty(mock_log.call_args_list) - def checkBinaryInputs(self, dataset_fn): + def _testBinaryInputs(self, dataset_fn): input1 = dataset_ops.Dataset.range(0) input2 = dataset_ops.Dataset.range(1) self.assertEqual([input1, input2], dataset_fn(input1, input2)._inputs()) @combinations.generate(test_base.default_test_combinations()) def testConcatenateInputs(self): - self.checkBinaryInputs(lambda x, y: x.concatenate(y)) + self._testBinaryInputs(lambda x, y: x.concatenate(y)) - def checkVariadicInputs(self, dataset_fn, input_datasets): + def _testVariadicInputs(self, dataset_fn, input_datasets): self.assertEqual( nest.flatten(input_datasets), dataset_fn(input_datasets)._inputs()) @@ -235,20 +234,20 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) def testZipOneInputs(self): input_datasets = dataset_ops.Dataset.range(0) - self.checkVariadicInputs(dataset_ops.Dataset.zip, input_datasets) + self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets) @combinations.generate(test_base.default_test_combinations()) def testZipNestInputs(self): input_datasets = (dataset_ops.Dataset.range(0), (dataset_ops.Dataset.range(1), dataset_ops.Dataset.range(2))) - self.checkVariadicInputs(dataset_ops.Dataset.zip, input_datasets) + self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets) @combinations.generate(test_base.default_test_combinations()) def testZipTupleInputs(self): input_datasets = (dataset_ops.Dataset.range(0), dataset_ops.Dataset.range(1)) - self.checkVariadicInputs(dataset_ops.Dataset.zip, input_datasets) + self._testVariadicInputs(dataset_ops.Dataset.zip, input_datasets) @combinations.generate(test_base.default_test_combinations()) def testFunctions(self): @@ -273,7 +272,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual(2, inputs.count(ds2)) self.assertEqual(1, inputs.count(ds3)) - def checkDatasetSpec(self, tf_value, expected_element_structure): + def _testDatasetSpec(self, tf_value, expected_element_structure): dataset = dataset_ops.Dataset.from_tensors(0).map(lambda _: tf_value) dataset_structure = structure.type_spec_from_value(dataset) self.assertIsInstance(dataset_structure, dataset_ops.DatasetSpec) @@ -307,12 +306,12 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) def testTensorDatasetSpec(self): - self.checkDatasetSpec( + self._testDatasetSpec( constant_op.constant(37.0), tensor_spec.TensorSpec([], dtypes.float32)) @combinations.generate(test_base.default_test_combinations()) def testSparseTensorDatasetSpec(self): - self.checkDatasetSpec( + self._testDatasetSpec( sparse_tensor.SparseTensor( indices=[[0]], values=constant_op.constant([0], dtype=dtypes.int32), @@ -320,7 +319,7 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) def testNestDatasetSpec(self): - self.checkDatasetSpec( + self._testDatasetSpec( { "a": constant_op.constant(37.0), "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar")) @@ -335,20 +334,19 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) def testDatasetDatasetSpec(self): - self.checkDatasetSpec( + self._testDatasetSpec( dataset_ops.Dataset.from_tensor_slices( constant_op.constant([1, 2, 3])), dataset_ops.DatasetSpec(tensor_spec.TensorSpec([], dtypes.int32))) @combinations.generate(test_base.default_test_combinations()) def testOptionalDatasetSpec(self): - self.checkDatasetSpec( + self._testDatasetSpec( optional_ops.Optional.from_value(37.0), optional_ops.OptionalSpec(tensor_spec.TensorSpec([], dtypes.float32))) - @combinations.generate( - combinations.combine(tf_api_version=[1], mode=["graph"])) - def testSkipEagerSameGraphErrorOneShot(self): + @combinations.generate(test_base.graph_only_combinations()) + def testSameGraphError(self): dataset = dataset_ops.Dataset.range(10) with ops.Graph().as_default(): with self.assertRaisesRegexp(ValueError, "must be from the same graph"): @@ -356,26 +354,27 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate( combinations.combine(tf_api_version=[1], mode=["graph"])) - def testSkipEagerSameGraphErrorOneShotSimple(self): + def testSameGraphErrorOneShot(self): dataset = dataset_ops.Dataset.range(10) with ops.Graph().as_default(): - with test.mock.patch.object(tf_logging, "warning") as mock_log: + with self.assertRaisesRegexp( + ValueError, "Please ensure that all datasets in the pipeline are " + "created in the same graph as the iterator."): _ = dataset_ops.make_one_shot_iterator(dataset) - self.assertRegexpMatches( - str(mock_log.call_args), "Please ensure that all datasets in the " - "pipeline are created in the same graph as the iterator.") @combinations.generate( combinations.combine(tf_api_version=[1], mode=["graph"])) - def testSkipEagerSameGraphErrorInitializable(self): + def testSameGraphErrorInitializable(self): dataset = dataset_ops.Dataset.range(10) with ops.Graph().as_default(): - with self.assertRaisesRegexp(ValueError, "must be from the same graph"): - dataset = dataset.batch(2) + with self.assertRaisesRegexp( + ValueError, "Please ensure that all datasets in the pipeline are " + "created in the same graph as the iterator."): + _ = dataset_ops.make_initializable_iterator(dataset) @combinations.generate( combinations.times( - combinations.combine(tf_api_version=[1, 2], mode="eager"), + test_base.eager_only_combinations(), combinations.combine(execution_mode=[context.ASYNC, context.SYNC]))) def testEagerIteration(self, execution_mode): with context.execution_mode(execution_mode): diff --git a/tensorflow/python/data/kernel_tests/filter_test.py b/tensorflow/python/data/kernel_tests/filter_test.py index 05b538a46ce..f6bdcb12020 100644 --- a/tensorflow/python/data/kernel_tests/filter_test.py +++ b/tensorflow/python/data/kernel_tests/filter_test.py @@ -30,28 +30,31 @@ from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -def new_and_legacy_filter_fn_combinations(): +def _test_combinations(): - def new_filter_fn(dataset, predicate): + def filter_fn(dataset, predicate): return dataset.filter(predicate) def legacy_filter_fn(dataset, predicate): return dataset.filter_with_legacy_function(predicate) - return (combinations.combine( + filter_combinations = combinations.combine( tf_api_version=[1, 2], mode=["eager", "graph"], - apply_filter=combinations.NamedObject("new_filter_fn", new_filter_fn)) + - combinations.combine( - tf_api_version=1, - mode=["eager", "graph"], - apply_filter=combinations.NamedObject("legacy_filter_fn", - legacy_filter_fn))) + apply_filter=combinations.NamedObject("filter_fn", filter_fn)) + + legacy_filter_combinations = combinations.combine( + tf_api_version=1, + mode=["eager", "graph"], + apply_filter=combinations.NamedObject("legacy_filter_fn", + legacy_filter_fn)) + + return filter_combinations + legacy_filter_combinations class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testFilterDataset(self, apply_filter): components = (np.arange(7, dtype=np.int64), np.array([[1, 2, 3]], dtype=np.int64) * @@ -87,14 +90,14 @@ class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): # Test an empty dataset. do_test(0, 1) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testFilterRange(self, apply_filter): dataset = dataset_ops.Dataset.range(4) dataset = apply_filter(dataset, lambda x: math_ops.not_equal(math_ops.mod(x, 3), 2)) self.assertDatasetProduces(dataset, expected_output=[0, 1, 3]) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testFilterDict(self, apply_filter): dataset = dataset_ops.Dataset.range(10).map( lambda x: {"foo": x * 2, "bar": x**2}) @@ -104,7 +107,7 @@ class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): dataset, expected_output=[(i * 2 + i**2) for i in range(10) if not (i**2) % 2]) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testUseStepContainerInFilter(self, apply_filter): input_data = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64) @@ -119,7 +122,7 @@ class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): dataset = apply_filter(dataset, _predicate) self.assertDatasetProduces(dataset, expected_output=[input_data[0]]) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testSparse(self, apply_filter): def _map_fn(i): @@ -137,7 +140,7 @@ class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( dataset, expected_output=[_map_fn(i * 2)[0] for i in range(5)]) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testShortCircuit(self, apply_filter): dataset = dataset_ops.Dataset.zip( (dataset_ops.Dataset.range(10), @@ -146,7 +149,7 @@ class FilterTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( dataset, expected_output=[(i, True) for i in range(10)]) - @combinations.generate(new_and_legacy_filter_fn_combinations()) + @combinations.generate(_test_combinations()) def testParallelFilters(self, apply_filter): dataset = dataset_ops.Dataset.range(10) dataset = apply_filter(dataset, lambda x: math_ops.equal(x % 2, 0)) diff --git a/tensorflow/python/data/kernel_tests/flat_map_test.py b/tensorflow/python/data/kernel_tests/flat_map_test.py index 3afed61fc7f..00b6e400ea7 100644 --- a/tensorflow/python/data/kernel_tests/flat_map_test.py +++ b/tensorflow/python/data/kernel_tests/flat_map_test.py @@ -66,10 +66,8 @@ class FlatMapTest(test_base.DatasetTestBase, parameterized.TestCase): expected_output.extend([i] * i) self.assertDatasetProduces(dataset, expected_output=expected_output) - # Note: no eager mode coverage, session specific test. - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) - def testSkipEagerSharedResourceNestedFlatMapDataset(self): + @combinations.generate(test_base.graph_only_combinations()) + def testSharedResourceNestedFlatMapDataset(self): repeats = [[1, 2], [3, 4], [5, 0], [1, 7]] components = np.array(repeats, dtype=np.int64) iterator = ( diff --git a/tensorflow/python/data/kernel_tests/from_generator_test.py b/tensorflow/python/data/kernel_tests/from_generator_test.py index dfa467add62..49753babacb 100644 --- a/tensorflow/python/data/kernel_tests/from_generator_test.py +++ b/tensorflow/python/data/kernel_tests/from_generator_test.py @@ -32,62 +32,83 @@ from tensorflow.python.ops import script_ops from tensorflow.python.platform import test -class DatasetConstructorTest(test_base.DatasetTestBase, parameterized.TestCase): +class FromGeneratorTest(test_base.DatasetTestBase, parameterized.TestCase): def _testFromGenerator(self, generator, elem_sequence, num_repeats, - output_types=None): - if output_types is None: - output_types = dtypes.int64 - dataset = dataset_ops.Dataset.from_generator( - generator, output_types=output_types).repeat(num_repeats).prefetch(5) - self.assertDatasetProduces( - dataset, - elem_sequence * num_repeats, - requires_initialization=True, - num_test_iterations=2) - - def _testFromGeneratorOneShot(self, generator, elem_sequence, num_repeats): + requires_initialization): dataset = dataset_ops.Dataset.from_generator( generator, output_types=dtypes.int64).repeat(num_repeats).prefetch(5) self.assertDatasetProduces( - dataset, elem_sequence * num_repeats, num_test_iterations=2) + dataset, + elem_sequence * num_repeats, + requires_initialization=requires_initialization, + num_test_iterations=2) + + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + num_repeats=[1, 5], requires_initialization=[True, False]))) + def testFromGeneratorUsingFn(self, num_repeats, requires_initialization): - @combinations.generate(test_base.default_test_combinations()) - def testFromGeneratorUsingFunction(self): def generator(): for i in range(1, 100): yield [i] * i - elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) - self._testFromGeneratorOneShot(generator, elem_sequence, 1) - self._testFromGeneratorOneShot(generator, elem_sequence, 5) - @combinations.generate(test_base.default_test_combinations()) - def testFromGeneratorUsingList(self): + elem_sequence = list(generator()) + self._testFromGenerator( + generator, + elem_sequence, + num_repeats=num_repeats, + requires_initialization=requires_initialization) + + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + num_repeats=[1, 5], requires_initialization=[True, False]))) + def testFromGeneratorUsingList(self, num_repeats, requires_initialization): generator = lambda: [[i] * i for i in range(1, 100)] elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) + self._testFromGenerator( + generator, + elem_sequence, + num_repeats=num_repeats, + requires_initialization=requires_initialization) - @combinations.generate(test_base.default_test_combinations()) - def testFromGeneratorUsingNdarray(self): + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + num_repeats=[1, 5], requires_initialization=[True, False]))) + def testFromGeneratorUsingNdarray(self, num_repeats, requires_initialization): generator = lambda: np.arange(100, dtype=np.int64) elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1, output_types=np.int64) - self._testFromGenerator(generator, elem_sequence, 5, output_types=np.int64) + self._testFromGenerator( + generator, + elem_sequence, + num_repeats=num_repeats, + requires_initialization=requires_initialization) - @combinations.generate(test_base.default_test_combinations()) - def testFromGeneratorUsingGeneratorExpression(self): - # NOTE(mrry): Generator *expressions* are not repeatable (or in - # general reusable), because they eagerly evaluate the `for` - # expression as `iter(range(1, 100))` and discard the means of - # reconstructing `range(1, 100)`. Wrapping the generator - # expression in a `lambda` makes it repeatable. + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + num_repeats=[1, 5], requires_initialization=[True, False]))) + def testFromGeneratorUsingGeneratorExpression(self, num_repeats, + requires_initialization): + # NOTE(mrry): Generator *expressions* are not repeatable (or in general + # reusable), because they eagerly evaluate the `for` expression as + # `iter(range(1, 100))` and discard the means of reconstructing + # `range(1, 100)`. Wrapping the generator expression in a `lambda` makes + # it repeatable. generator = lambda: ([i] * i for i in range(1, 100)) elem_sequence = list(generator()) - self._testFromGenerator(generator, elem_sequence, 1) - self._testFromGenerator(generator, elem_sequence, 5) + self._testFromGenerator( + generator, + elem_sequence, + num_repeats=num_repeats, + requires_initialization=requires_initialization) @combinations.generate(test_base.default_test_combinations()) def testFromMultipleConcurrentGenerators(self): @@ -392,7 +413,6 @@ class DatasetConstructorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertAllEqual(37, self.evaluate(get_next())) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - self.assertTrue(event.is_set()) @combinations.generate(test_base.default_test_combinations()) def testSharedName(self): diff --git a/tensorflow/python/data/kernel_tests/from_tensors_test.py b/tensorflow/python/data/kernel_tests/from_tensors_test.py index e293383403f..c899c156739 100644 --- a/tensorflow/python/data/kernel_tests/from_tensors_test.py +++ b/tensorflow/python/data/kernel_tests/from_tensors_test.py @@ -237,8 +237,8 @@ class FromTensorsTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual([3], get_next().shape) # TODO(b/121264236): needs mechanism for multiple device in eager mode. - @combinations.generate(test_base.default_test_combinations()) - def testSkipEagerSplitPipeline(self): + @combinations.generate(test_base.graph_only_combinations()) + def testSplitPipeline(self): with session.Session( target="", config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess: diff --git a/tensorflow/python/data/kernel_tests/iterator_cluster_test.py b/tensorflow/python/data/kernel_tests/iterator_cluster_test.py index 0384f9fc18a..0a40c212006 100644 --- a/tensorflow/python/data/kernel_tests/iterator_cluster_test.py +++ b/tensorflow/python/data/kernel_tests/iterator_cluster_test.py @@ -17,12 +17,15 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized import numpy as np from tensorflow.core.protobuf import config_pb2 from tensorflow.python.client import session +from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import iterator_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -37,9 +40,9 @@ from tensorflow.python.ops import string_ops from tensorflow.python.platform import test -class IteratorClusterTest(test.TestCase): +class IteratorClusterTest(test.TestCase, parameterized.TestCase): - @test_util.run_v1_only("b/120545219") + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorWithoutRemoteCallFail(self): worker_config = config_pb2.ConfigProto() worker_config.device_count["CPU"] = 2 @@ -95,7 +98,7 @@ class IteratorClusterTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(remote_op, feed_dict={target_placeholder: device1}) - @test_util.run_v1_only("b/120545219") + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorUsingRemoteCallOp(self): worker_config = config_pb2.ConfigProto() worker_config.device_count["CPU"] = 2 @@ -106,7 +109,7 @@ class IteratorClusterTest(test.TestCase): "/job:worker/replica:0/task:0/cpu:1", worker[0].target) - @test_util.run_v1_only("b/120545219") + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorUsingRemoteCallOpCrossProcess(self): workers, _ = test_util.create_local_cluster(2, 1) @@ -114,7 +117,7 @@ class IteratorClusterTest(test.TestCase): "/job:worker/replica:0/task:1/cpu:0", workers[0].target) - @test_util.run_v1_only("b/120545219") + @combinations.generate(test_base.graph_only_combinations()) def testCaptureHashTableInSharedIterator(self): worker, _ = test_util.create_local_cluster(1, 1) @@ -131,10 +134,10 @@ class IteratorClusterTest(test.TestCase): input_sentences = dataset_ops.Dataset.from_tensor_slices( ["brain brain tank salad surgery", "surgery brain"]) - iterator = ( - input_sentences.map(lambda x: string_ops.string_split([x]).values).map( - table.lookup) - .make_initializable_iterator(shared_name="shared_iterator")) + dataset = input_sentences.map( + lambda x: string_ops.string_split([x]).values).map(table.lookup) + iterator = dataset_ops.make_initializable_iterator( + dataset, shared_name="shared_iterator") init_op = iterator.initializer get_next = iterator.get_next() @@ -148,7 +151,7 @@ class IteratorClusterTest(test.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @test_util.run_v1_only("b/120545219") + @combinations.generate(test_base.graph_only_combinations()) def testImplicitDisposeParallelMapDataset(self): # Tests whether a parallel map dataset will be cleaned up correctly when # the pipeline does not run it until exhaustion. diff --git a/tensorflow/python/data/kernel_tests/iterator_test.py b/tensorflow/python/data/kernel_tests/iterator_test.py index 70f03f3e4d2..fcb2e4c0b1f 100644 --- a/tensorflow/python/data/kernel_tests/iterator_test.py +++ b/tensorflow/python/data/kernel_tests/iterator_test.py @@ -56,8 +56,7 @@ from tensorflow.python.util import compat class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testNoGradients(self): component = constant_op.constant([1.]) side = constant_op.constant(0.) @@ -68,8 +67,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertIsNone(gradients_impl.gradients(value, side)[0]) self.assertIsNone(gradients_impl.gradients(value, [component, side])[0]) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testCapturingStateInOneShotRaisesException(self): var = variables.Variable(37.0, name="myvar") dataset = ( @@ -80,8 +78,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): "datasets that capture stateful objects.+myvar"): dataset_ops.make_one_shot_iterator(dataset) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testOneShotIterator(self): components = (np.arange(7), np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], @@ -107,8 +104,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testOneShotIteratorCaptureByValue(self): components = (np.arange(7), np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], @@ -172,8 +168,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testOneShotIteratorNonBlocking(self): dataset = dataset_ops.Dataset.from_tensors([1, 2, 3]).map(lambda x: x * x) iterator = dataset_ops.make_one_shot_iterator(dataset) @@ -207,13 +202,11 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): for t in threads: t.join() - self.assertEqual(num_threads, len(results)) - self.assertEqual(num_threads - 1, - len([None for r in results if r is None])) + self.assertLen(results, num_threads) + self.assertLen([None for r in results if r is None], num_threads - 1) self.assertAllEqual([[1, 4, 9]], [r for r in results if r is not None]) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testOneShotIteratorInitializerFails(self): # Define a dataset whose initialization will always fail. dataset = dataset_ops.Dataset.from_tensors(array_ops.gather([0], [4])) @@ -243,8 +236,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): for t in threads: t.join() - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testSimpleSharedResource(self): components = (np.array(1, dtype=np.int64), np.array([1, 2, 3], dtype=np.int64), @@ -294,8 +286,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testNotInitializedError(self): components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) iterator = dataset_ops.make_initializable_iterator( @@ -307,8 +298,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): "iterator has not been initialized"): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testReinitializableIterator(self): dataset_3 = dataset_ops.Dataset.from_tensors( constant_op.constant([1, 2, 3])) @@ -353,8 +343,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(get_next) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testReinitializableIteratorWithFunctions(self): def g(): @@ -415,8 +404,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): (constant_op.constant([1, 2, 3], dtype=dtypes.int64), constant_op.constant([4., 5., 6., 7.], dtype=dtypes.float64)))) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testIteratorStringHandle(self): dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) dataset_4 = dataset_ops.Dataset.from_tensor_slices([10, 20, 30, 40]) @@ -474,8 +462,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run( next_element, feed_dict={handle_placeholder: iterator_4_handle}) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testIteratorStringHandleFuture(self): with forward_compat.forward_compatibility_horizon(2018, 8, 4): dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) @@ -541,8 +528,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): sess.run( next_element, feed_dict={handle_placeholder: iterator_4_handle}) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testIteratorStringHandleReuseTensorObject(self): dataset = dataset_ops.Dataset.from_tensor_slices([1, 2, 3]) one_shot_iterator = dataset_ops.make_one_shot_iterator(dataset) @@ -571,8 +557,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual("foo_1", handle_with_same_name.op.name) self.assertIsNot(handle_with_name, handle_with_same_name) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testIteratorStringHandleError(self): dataset_int_scalar = ( dataset_ops.Dataset.from_tensor_slices([1, 2, 3]).repeat()) @@ -613,8 +598,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): feedable_int_vector.get_next(), feed_dict={handle_placeholder: handle_float_vector})) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorUsingRemoteCallOpDirectSession(self): worker_config = config_pb2.ConfigProto() worker_config.device_count["CPU"] = 3 @@ -672,8 +656,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): target_placeholder: "/job:localhost/replica:0/task:0/cpu:1" }) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorUsingRemoteCallOpMultiWorkers(self): s1 = server_lib.Server.create_local_server() s2 = server_lib.Server.create_local_server() @@ -727,8 +710,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(n) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testRemoteIteratorUsingRemoteCallOpDirectSessionGPUCPU(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") @@ -785,8 +767,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): target_placeholder: "/job:localhost/replica:0/task:0/cpu:0" }) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode=["graph"])) + @combinations.generate(test_base.graph_only_combinations()) def testRepeatedGetNextWarning(self): iterator = dataset_ops.make_one_shot_iterator(dataset_ops.Dataset.range(10)) warnings.simplefilter("always") @@ -929,7 +910,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual(val, foo.numpy()) val += 1 - @combinations.generate(combinations.combine(tf_api_version=2, mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testOwnedIteratorFunction(self): queue = data_flow_ops.FIFOQueue(10, dtypes.int64) @@ -946,7 +927,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): for i in range(10): self.assertEqual(queue.dequeue().numpy(), i) - @combinations.generate(combinations.combine(tf_api_version=2, mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testOwnedIteratorFunctionError(self): # In this test we verify that a function that raises an error ends up # properly deallocating the iterator resource. @@ -976,7 +957,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual(queue.size().numpy(), 2) - @combinations.generate(combinations.combine(tf_api_version=2, mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testLimitedRetracing(self): trace_count = [0] @@ -996,7 +977,7 @@ class IteratorTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual(self.evaluate(f(iter(dataset2))), 45) self.assertEqual(trace_count[0], 1) - @combinations.generate(combinations.combine(tf_api_version=2, mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testNestedFunctionsIteratorResource(self): @def_function.function diff --git a/tensorflow/python/data/kernel_tests/list_files_test.py b/tensorflow/python/data/kernel_tests/list_files_test.py index 52ce300f537..40b4b77116c 100644 --- a/tensorflow/python/data/kernel_tests/list_files_test.py +++ b/tensorflow/python/data/kernel_tests/list_files_test.py @@ -35,10 +35,12 @@ from tensorflow.python.util import compat class ListFilesTest(test_base.DatasetTestBase, parameterized.TestCase): def setUp(self): + super(ListFilesTest, self).setUp() self.tmp_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.tmp_dir, ignore_errors=True) + super(ListFilesTest, self).tearDown() def _touchTempFiles(self, filenames): for filename in filenames: diff --git a/tensorflow/python/data/kernel_tests/map_test.py b/tensorflow/python/data/kernel_tests/map_test.py index 0847cdd7a0d..c8b23edbc7f 100644 --- a/tensorflow/python/data/kernel_tests/map_test.py +++ b/tensorflow/python/data/kernel_tests/map_test.py @@ -17,6 +17,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools from collections import namedtuple import threading import time @@ -31,13 +32,13 @@ from tensorflow.python.data.experimental.ops import threading_options from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_util -from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import data_flow_ops @@ -57,10 +58,70 @@ from tensorflow.python.ops.ragged import ragged_tensor from tensorflow.python.platform import test -def _make_coordinated_sloppy_dataset(num_elements, num_parallel_calls): +def _test_combinations_with_mode_v1(mode): + + def new_map_fn(dataset, *args, **kwargs): + return dataset.map(*args, **kwargs) + + def legacy_map_fn(dataset, *args, **kwargs): + return dataset.map_with_legacy_function(*args, **kwargs) + + new_map_combinations = combinations.combine( + tf_api_version=1, + mode=mode, + apply_map=combinations.NamedObject("map_fn", new_map_fn)) + + legacy_map_combinations = combinations.combine( + tf_api_version=1, + mode=mode, + apply_map=combinations.NamedObject("legacy_map_fn", legacy_map_fn)) + + return new_map_combinations + legacy_map_combinations + + +def _test_combinations_with_mode_v2(mode): + + def new_map_fn(dataset, *args, **kwargs): + return dataset.map(*args, **kwargs) + + return combinations.combine( + tf_api_version=2, + mode=mode, + apply_map=combinations.NamedObject("map_fn", new_map_fn)) + + +def _test_combinations_with_mode(mode): + return _test_combinations_with_mode_v1( + mode) + _test_combinations_with_mode_v2(mode) + + +def _test_combinations(): + return _test_combinations_with_mode("eager") + _test_combinations_with_mode( + "graph") + + +def _short_circuit_test_cases(): + cases = [ + ("Identity", None, lambda x: x), + ("Replicate", None, lambda x: (x, x)), + ("Swap", (None, None), lambda x, y: (y, x)), + ("Project", (None, None), lambda x, y: x) + ] + + def reduce_fn(x, y): + name, structure, fn = y + return x + combinations.combine( + structure=structure, fn=combinations.NamedObject(name, fn)) + + return functools.reduce(reduce_fn, cases, []) + + +def _make_coordinated_sloppy_dataset(apply_map, num_elements, + num_parallel_calls): """Produces a dataset iterator and events to control the order of elements. Args: + apply_map: method that applies the `map` transformation num_elements: the number of input elements num_parallel_calls: the degree of map parallelism @@ -84,28 +145,27 @@ def _make_coordinated_sloppy_dataset(num_elements, num_parallel_calls): options = dataset_ops.Options() options.experimental_deterministic = False - dataset = dataset_ops.Dataset.range(num_elements).map( - map_fn, num_parallel_calls).with_options(options) + dataset = dataset_ops.Dataset.range(num_elements) + dataset = apply_map(dataset, map_fn, num_parallel_calls).with_options(options) return dataset, coordination_events -# TODO(jsimsa): Add tests for `map_with_legacy_function`. -@test_util.run_all_in_graph_and_eager_modes class MapTest(test_base.DatasetTestBase, parameterized.TestCase): - def _buildMapDataset(self, components, count): + def _map_dataset_factory(self, components, apply_map, count): def _map_fn(x, y, z): return math_ops.square(x), math_ops.square(y), math_ops.square(z) - dataset = dataset_ops.Dataset.from_tensor_slices(components).map( - _map_fn).repeat(count) + dataset = dataset_ops.Dataset.from_tensor_slices(components) + dataset = apply_map(dataset, _map_fn).repeat(count) self.assertEqual( [c.shape[1:] for c in components], [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) return dataset - def testMapDataset(self): + @combinations.generate(_test_combinations()) + def testMapDataset(self, apply_map): """Test an dataset that maps a TF function across its input elements.""" # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> # RepeatDataset(count). @@ -114,7 +174,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): np.array(37.0) * np.arange(7)) # Test single-threaded access to the iterator. - get_next = self.getNext(self._buildMapDataset(components, 14)) + get_next = self.getNext( + self._map_dataset_factory(components, apply_map, count=14)) for _ in range(14): for i in range(7): result = self.evaluate(get_next()) @@ -123,15 +184,15 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - # TODO(b/117581999): add eager coverage, different threads run in graph - # context. - @test_util.run_v1_only("b/120545219") - def testSkipEagerMapDatasetMultithreaded(self): + # TODO(b/117581999): add eager coverage + @combinations.generate(_test_combinations_with_mode("graph")) + def testMapDatasetMultiThreaded(self, apply_map): # Test multi-threaded access to the same iterator. components = (np.arange(7), np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], np.array(37.0) * np.arange(7)) - get_next = self.getNext(self._buildMapDataset(components, 18)) + get_next = self.getNext( + self._map_dataset_factory(components, apply_map, count=18)) results = [] with self.cached_session() as sess: def iterator_thread(): @@ -157,94 +218,99 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): results[i * 18 + j]): self.assertAllEqual(component[i]**2, result_component) - def _buildParallelMapDataset(self, components, count, num_parallel_calls, - output_buffer_size): + def _parallel_map_dataset_factory(self, components, apply_map, count, + num_parallel_calls, buffer_size): def _map_fn(x, y, z): return math_ops.square(x), math_ops.square(y), math_ops.square(z) - dataset = dataset_ops.Dataset.from_tensor_slices(components).map( - _map_fn, num_parallel_calls=num_parallel_calls).prefetch( - output_buffer_size).repeat(count) + dataset = dataset_ops.Dataset.from_tensor_slices(components) + dataset = apply_map(dataset, _map_fn, num_parallel_calls=num_parallel_calls) + dataset = dataset.prefetch(buffer_size).repeat(count) self.assertEqual( [c.shape[1:] for c in components], [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) return dataset - def testParallelMapDataset(self): + @combinations.generate( + combinations.times( + _test_combinations(), + combinations.combine(num_parallel_calls=1, buffer_size=1) + + combinations.combine(num_parallel_calls=1, buffer_size=2) + + combinations.combine(num_parallel_calls=2, buffer_size=2) + + combinations.combine(num_parallel_calls=2, buffer_size=4) + + combinations.combine(num_parallel_calls=8, buffer_size=8) + + combinations.combine(num_parallel_calls=8, buffer_size=16))) + def testParallelMapDataset(self, apply_map, num_parallel_calls, buffer_size): """Test an dataset that maps a TF function across its input elements.""" # The pipeline is TensorSliceDataset -> ParallelMapDataset(square_3) -> # RepeatDataset(count). - def do_test(num_parallel_calls, output_buffer_size): + components = (np.arange(7), + np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], + np.array(37.0) * np.arange(7)) + # Test single-threaded access to the iterator. + get_next = self.getNext( + self._parallel_map_dataset_factory(components, apply_map, 14, + num_parallel_calls, buffer_size)) + for _ in range(14): + for i in range(7): + result = self.evaluate(get_next()) + for component, result_component in zip(components, result): + self.assertAllEqual(component[i]**2, result_component) + with self.assertRaises(errors.OutOfRangeError): + self.evaluate(get_next()) - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - # Test single-threaded access to the iterator. - get_next = self.getNext( - self._buildParallelMapDataset(components, 14, num_parallel_calls, - output_buffer_size)) - for _ in range(14): - for i in range(7): - result = self.evaluate(get_next()) - for component, result_component in zip(components, result): + # TODO(b/117581999): add eager coverage + @combinations.generate( + combinations.times( + _test_combinations_with_mode("graph"), + combinations.combine(num_parallel_calls=1, buffer_size=1) + + combinations.combine(num_parallel_calls=1, buffer_size=2) + + combinations.combine(num_parallel_calls=2, buffer_size=2) + + combinations.combine(num_parallel_calls=2, buffer_size=4) + + combinations.combine(num_parallel_calls=8, buffer_size=8) + + combinations.combine(num_parallel_calls=8, buffer_size=16))) + def testParallelMapDatasetMultiThreaded(self, apply_map, num_parallel_calls, + buffer_size): + + # Test multi-threaded access to the same iterator. + components = (np.arange(7), + np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], + np.array(37.0) * np.arange(7)) + get_next = self.getNext( + self._parallel_map_dataset_factory(components, apply_map, 18, + num_parallel_calls, buffer_size)) + results = [] + with self.cached_session() as sess: + + def iterator_thread(): + while True: + try: + results.append(sess.run(get_next())) + except errors.OutOfRangeError: + return + + threads = [self.checkedThread(target=iterator_thread) for _ in range(64)] + for t in threads: + t.start() + for t in threads: + t.join() + + # `results` will contain the same elements components**2 + # repeated 18 times, but in a non-deterministic order. Sort the + # results, and assert that each element of components**2 is + # produced 18 times. + results.sort(key=lambda x: x[0]) + for i in range(7): + for j in range(18): + for component, result_component in zip(components, + results[i * 18 + j]): self.assertAllEqual(component[i]**2, result_component) - with self.assertRaises(errors.OutOfRangeError): - self.evaluate(get_next()) - for num_parallel_calls_val, output_buffer_size_val in [(1, 1), (1, 2), (2, - 2), - (2, 4), (8, 8), - (8, 16)]: - do_test(num_parallel_calls_val, output_buffer_size_val) - - # TODO(b/117581999): add eager coverage, different threads run in graph - # context. - @test_util.run_v1_only("b/120545219") - def testSkipEagerParallelMapDatasetMultithreaded(self): - - def do_test(num_parallel_calls, output_buffer_size): - # Test multi-threaded access to the same iterator. - components = (np.arange(7), - np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], - np.array(37.0) * np.arange(7)) - get_next = self.getNext( - self._buildParallelMapDataset(components, 18, num_parallel_calls, - output_buffer_size)) - results = [] - with self.cached_session() as sess: - - def iterator_thread(): - while True: - try: - results.append(sess.run(get_next())) - except errors.OutOfRangeError: - return - threads = [self.checkedThread(target=iterator_thread) - for _ in range(64)] - for t in threads: - t.start() - for t in threads: - t.join() - - # `results` will contain the same elements components**2 - # repeated 18 times, but in a non-deterministic order. Sort the - # results, and assert that each element of components**2 is - # produced 18 times. - results.sort(key=lambda x: x[0]) - for i in range(7): - for j in range(18): - for component, result_component in zip(components, - results[i * 18 + j]): - self.assertAllEqual(component[i]**2, result_component) - - for num_parallel_calls_val, output_buffer_size_val in [ - (1, 1), (1, 2), (2, 2), (2, 4), (8, 8), (8, 16)]: - do_test(num_parallel_calls_val, output_buffer_size_val) - - def testImplicitDisposeParallelMapDataset(self): + @combinations.generate(_test_combinations()) + def testImplicitDisposeParallelMapDataset(self, apply_map): # Tests whether a parallel map dataset will be cleaned up correctly when # the pipeline does not run it until exhaustion. # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> @@ -253,7 +319,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): np.array([[1, 2, 3]]) * np.arange(1000)[:, np.newaxis], np.array(37.0) * np.arange(1000)) - dataset = self._buildParallelMapDataset(components, 1000, 100, 100) + dataset = self._parallel_map_dataset_factory(components, apply_map, 1000, + 100, 100) # NOTE(mrry): Also test that the prefetching thread is cancelled correctly. dataset = dataset.prefetch(100) get_next = self.getNext(dataset) @@ -261,23 +328,29 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): for _ in range(3): self.evaluate(get_next()) - def testParallelMapUnspecifiedOutputSize(self): + @combinations.generate(_test_combinations()) + def testParallelMapUnspecifiedOutputSize(self, apply_map): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message"), - num_parallel_calls=2)) + dataset = dataset_ops.Dataset.from_tensor_slices(components) + dataset = apply_map( + dataset, + lambda x: array_ops.check_numerics(x, "message"), + num_parallel_calls=2) get_next = self.getNext(dataset) for _ in range(3): self.evaluate(get_next()) - def testParallelMapError(self): + @combinations.generate(_test_combinations()) + def testParallelMapError(self, apply_map): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message"), - num_parallel_calls=2)) + dataset = dataset_ops.Dataset.from_tensor_slices(components) + dataset = apply_map( + dataset, + lambda x: array_ops.check_numerics(x, "message"), + num_parallel_calls=2) get_next = self.getNext(dataset) for _ in range(3): @@ -289,13 +362,13 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testPrefetchError(self): + @combinations.generate(_test_combinations()) + def testPrefetchError(self, apply_map): components = np.array([1., 2., 3., np.nan, 5.]).astype(np.float32) - dataset = (dataset_ops.Dataset.from_tensor_slices(components) - .map(lambda x: array_ops.check_numerics(x, "message")) - .prefetch(2)) - + dataset = dataset_ops.Dataset.from_tensor_slices(components) + dataset = apply_map( + dataset, lambda x: array_ops.check_numerics(x, "message")).prefetch(2) get_next = self.getNext(dataset) for _ in range(3): @@ -307,7 +380,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testCaptureIterator(self): + @combinations.generate(_test_combinations()) + def testCaptureIterator(self, apply_map): def _build_ds(iterator): @@ -315,7 +389,7 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): get_next = iterator.get_next() return x * get_next - return dataset_ops.Dataset.range(10).map(_map_fn) + return apply_map(dataset_ops.Dataset.range(10), _map_fn) def _build_graph(): if context.executing_eagerly(): @@ -335,7 +409,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testCaptureHashTable(self): + @combinations.generate(_test_combinations()) + def testCaptureHashTable(self, apply_map): # NOTE(mrry): We must use the V2 variants of `HashTable` # etc. because these produce a `tf.resource`-typed output that is # compatible with the in-graph function implementation. @@ -348,8 +423,9 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): input_sentences = dataset_ops.Dataset.from_tensor_slices( ["brain brain tank salad surgery", "surgery brain"]) - dataset = input_sentences.map(lambda x: string_ops.string_split([x]).values - ).map(table.lookup) + dataset = apply_map(input_sentences, + lambda x: string_ops.string_split([x]).values) + dataset = apply_map(dataset, table.lookup) get_next = self.getNext(dataset, requires_initialization=True) @@ -359,14 +435,15 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @test_util.run_v1_only("b/123904513") - def testCaptureQueue(self): + # TODO(b/123904513) + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testCaptureQueue(self, apply_map): elements = np.random.randint(100, size=[200]) queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[]) enqueue_op = queue.enqueue_many(elements) close_op = queue.close() - dataset = dataset_ops.Dataset.from_tensors(0).repeat( - -1).map(lambda _: queue.dequeue()) + dataset = dataset_ops.Dataset.from_tensors(0).repeat(-1) + dataset = apply_map(dataset, lambda _: queue.dequeue()) get_next = self.getNext(dataset, requires_initialization=True) self.evaluate(enqueue_op) @@ -378,8 +455,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.evaluate(get_next()) # TODO(b/117581999): Possible deadlock in eager mode, debug. - @test_util.run_v1_only("b/120545219") - def testSkipEagerCaptureSameResourceMultipleTimes(self): + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testCaptureSameResourceMultipleTimes(self, apply_map): elements = np.random.randint(100, size=[200]) queue = data_flow_ops.FIFOQueue( 200, dtypes.int64, shapes=[], shared_name="shared_queue") @@ -389,8 +466,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): enqueue_op = queue.enqueue_many(elements) close_op = queue.close() - dataset = dataset_ops.Dataset.from_tensors(0).repeat( - -1).map(lambda _: (queue.dequeue(), queue_2.dequeue())) + dataset = dataset_ops.Dataset.from_tensors(0).repeat(-1) + dataset = apply_map(dataset, lambda _: (queue.dequeue(), queue_2.dequeue())) self.evaluate(enqueue_op) self.evaluate(close_op) @@ -401,9 +478,11 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testSeededStatefulOperatorIsProperlyStateful(self): - dataset = dataset_ops.Dataset.from_tensors(0).repeat( - 10).map(lambda _: random_ops.random_uniform((), seed=11)).batch(2) + @combinations.generate(_test_combinations()) + def testSeededStatefulOperatorIsProperlyStateful(self, apply_map): + dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) + fn = lambda _: random_ops.random_uniform((), seed=11) + dataset = apply_map(dataset, fn).batch(2) get_next = self.getNext(dataset, requires_initialization=True) random_values = [] @@ -422,9 +501,11 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): # Randomness is repeatable given same seed self.assertAllClose(random_values, random_values_2) - def testStatefulMapKeepsStateAcrossIterators(self): - dataset = dataset_ops.Dataset.from_tensors(0).repeat(10).map( - lambda _: random_ops.random_uniform((), seed=11)).repeat(1000).batch(10) + @combinations.generate(_test_combinations()) + def testStatefulMapKeepsStateAcrossIterators(self, apply_map): + dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) + fn = lambda _: random_ops.random_uniform((), seed=11) + dataset = apply_map(dataset, fn).repeat(1000).batch(10) get_next = self.getNext(dataset) random_values = self.evaluate(get_next()) @@ -438,7 +519,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): i += 1 self.assertLess(i, 99) - def testStatefulOperationInShortCircuit(self): + @combinations.generate(_test_combinations()) + def testStatefulOperationInShortCircuit(self, apply_map): counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) @@ -446,7 +528,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): counter_var.assign_add(1) return x - dataset = dataset_ops.Dataset.range(10).map(increment_fn) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, increment_fn) get_next = self.getNext(dataset, requires_initialization=True) @@ -459,22 +542,24 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.evaluate(get_next()) self.assertEqual(10, self.evaluate(counter_var)) - def testMapDict(self): - dataset = dataset_ops.Dataset.range(10).map( - lambda x: {"foo": x * 2, "bar": x**2}).map( - lambda d: d["foo"] + d["bar"]) + @combinations.generate(_test_combinations()) + def testMapDict(self, apply_map): + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, lambda x: {"foo": x * 2, "bar": x**2}) + dataset = apply_map(dataset, lambda d: d["foo"] + d["bar"]) self.assertDatasetProduces( dataset, expected_output=[i * 2 + i**2 for i in range(10)]) - def testMapNamedtuple(self, count=10): + @combinations.generate(_test_combinations()) + def testMapNamedtuple(self, apply_map): # construct dataset of tuples - labels = dataset_ops.Dataset.range(count) - images = labels.map(lambda l: -l) + labels = dataset_ops.Dataset.range(10) + images = apply_map(labels, lambda l: -l) dataset_tuple = dataset_ops.Dataset.zip((labels, images)) # convert dataset of tuples to dataset of namedtuples example = namedtuple("Example", ["label", "image"]) - dataset_namedtuple = dataset_tuple.map(example) + dataset_namedtuple = apply_map(dataset_tuple, example) def preprocess_tuple(label, image): image = 2 * image @@ -484,14 +569,14 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): return example._replace(image=2 * example.image) # preprocess both datasets - dataset_tuple = dataset_tuple.map(preprocess_tuple) - dataset_namedtuple = dataset_namedtuple.map(preprocess_namedtuple) + dataset_tuple = apply_map(dataset_tuple, preprocess_tuple) + dataset_namedtuple = apply_map(dataset_namedtuple, preprocess_namedtuple) next_tuple = self.getNext(dataset_tuple) next_namedtuple = self.getNext(dataset_namedtuple) # make sure both datasets contain the same data - for i in range(count): + for i in range(10): tuple_, namedtuple_ = self.evaluate([next_tuple(), next_namedtuple()]) self.assertEqual(tuple_, namedtuple_) self.assertEqual(tuple_, (i, -2 * i)) @@ -499,13 +584,16 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_namedtuple()) - def testUseStepContainerInMap(self): + @combinations.generate(_test_combinations()) + def testUseStepContainerInMap(self, apply_map): row = np.arange(6) - dataset = dataset_ops.Dataset.from_tensors( - row).map(lambda elems: map_fn.map_fn(lambda x: x * x, elems)) + dataset = dataset_ops.Dataset.from_tensors(row) + dataset = apply_map(dataset, + lambda elems: map_fn.map_fn(lambda x: x * x, elems)) self.assertDatasetProduces(dataset, expected_output=[row**2]) - def testCaseAndCondInMap(self): + @combinations.generate(_test_combinations()) + def testCaseAndCondInMap(self, apply_map): def control_map_fn(x, y): @@ -531,13 +619,12 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): pred_fn_pairs, default=multiply, exclusive=True) def build_dataset(row, num): - dataset = dataset_ops.Dataset.from_tensor_slices( - row).map(lambda x: control_map_fn(x, num)) - return self.getNext(dataset) + dataset = dataset_ops.Dataset.from_tensor_slices(row) + return apply_map(dataset, lambda x: control_map_fn(x, num)) row = np.arange(6) for num in [2, 3, 4]: - get_next = build_dataset(row, num) + get_next = self.getNext(build_dataset(row, num)) for i in range(6): self.assertEqual( (i // 2 if i % 2 else i * 2) if (num == 2 or num == 3) else i * 2, @@ -545,7 +632,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testCaseInWhileInMap(self): + @combinations.generate(_test_combinations()) + def testCaseInWhileInMap(self, apply_map): def control_map_fn(x, y): @@ -564,22 +652,22 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): pred_fn_pairs, default=multiply, exclusive=True) def build_dataset(row, num): - # pylint: disable=g-long-lambda - dataset = dataset_ops.Dataset.from_tensors( - row).map(lambda elems: map_fn.map_fn( - lambda x: control_map_fn(x, num), elems)) - return self.getNext(dataset) + dataset = dataset_ops.Dataset.from_tensors(row) + return apply_map( + dataset, + lambda elems: map_fn.map_fn(lambda x: control_map_fn(x, num), elems)) row = np.arange(6) for num in [2, 3, 4]: - get_next = build_dataset(row, num) + get_next = self.getNext(build_dataset(row, num)) self.assertAllEqual( [x // 2 if (num == 2 or num == 3) else x * 2 for x in row], self.evaluate(get_next())) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testCaseAndCondInWhileInMap(self): + @combinations.generate(_test_combinations()) + def testCaseAndCondInWhileInMap(self, apply_map): def control_map_fn(x, y): @@ -606,11 +694,10 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): row = np.arange(6) num = 2 - # pylint: disable=g-long-lambda - dataset = dataset_ops.Dataset.from_tensors( - row).map(lambda elems: map_fn.map_fn( - lambda x: control_map_fn(x, num), elems)) - # pylint: enable=g-long-lambda + dataset = dataset_ops.Dataset.from_tensors(row) + dataset = apply_map( + dataset, + lambda elems: map_fn.map_fn(lambda x: control_map_fn(x, num), elems)) get_next = self.getNext(dataset) self.assertAllEqual([(x // 2 if x % 2 else x * 2) if @@ -619,17 +706,20 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testNestedListMapDataset(self): - dataset = dataset_ops.Dataset.from_tensors( - [0, 1, 2]).repeat(10).map(lambda a: ([a[1], a[0] + a[2]], a[1])) - + @combinations.generate(_test_combinations()) + def testNestedListMapDataset(self, apply_map): + dataset = dataset_ops.Dataset.from_tensors([0, 1, 2]).repeat(10) + dataset = apply_map(dataset, lambda a: ([a[1], a[0] + a[2]], a[1])) expected_output = [(np.array([1, 2]), 1)] * 10 self.assertDatasetProduces(dataset, expected_output=expected_output) - def testPrefetch(self): - # We will use this event to test that `_map_py_func()` has been - # invoked a certain number of times (6 times, to be exact) after - # consuming fewer elements from the iterator. + @combinations.generate( + combinations.times(_test_combinations(), + combinations.combine(buffer_size=[1, 2, 3, 4]))) + def testPrefetch(self, apply_map, buffer_size): + # We will use this event to test that `_map_py_func()` has been invoked a + # certain number of times (6 times, to be exact) after consuming fewer + # elements from the iterator. ev = threading.Event() set_event_during_invocation = 5 @@ -642,56 +732,38 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): def _map_fn(x): return script_ops.py_func(_map_py_func, [x], x.dtype) - def do_test(buffer_size): - dataset = dataset_ops.Dataset.range(100).map(_map_fn).prefetch( - buffer_size) - - get_next = self.getNext(dataset) - # Simple test that prefetch yields the expected values in the - # expected order. - for i in range(100): - self.assertEqual(i * i, self.evaluate(get_next())) - with self.assertRaises(errors.OutOfRangeError): - self.evaluate(get_next()) - - for buffer_size in [1, 10, 100, 1000]: - do_test(buffer_size) - - # We can indirectly observe that varying the buffer size has the - # intended effect by observing when `ev` is set (on the 6th - # invocation of `_map_py_func()`). + # We can indirectly observe that varying the buffer size has the intended + # effect by observing when `ev` is set (on the 6th invocation of + # `_map_py_func()`). # NOTE(mrry): We do not test with `buffer_size == - # set_event_during_invocation`, because we must consume at least - # one element to start the prefetching. - def do_test_ev(buffer_size): - dataset = dataset_ops.Dataset.range(100).map(_map_fn).prefetch( - buffer_size) + # set_event_during_invocation`, because we must consume at least one element + # to start the prefetching. + dataset = dataset_ops.Dataset.range(100) + dataset = apply_map(dataset, _map_fn).prefetch(buffer_size) + get_next = self.getNext(dataset) - get_next = self.getNext(dataset) + event_will_be_set_after_consuming = ( + set_event_during_invocation - buffer_size + 1) - event_will_be_set_after_consuming = ( - set_event_during_invocation - buffer_size + 1) + ev.clear() + for i in range(event_will_be_set_after_consuming): + self.assertFalse(ev.is_set()) + self.assertEqual(i * i, self.evaluate(get_next())) + ev.wait() + for i in range(event_will_be_set_after_consuming, 100): + self.assertEqual(i * i, self.evaluate(get_next())) + with self.assertRaises(errors.OutOfRangeError): + self.evaluate(get_next()) - ev.clear() - for i in range(event_will_be_set_after_consuming): - self.assertFalse(ev.is_set()) - self.assertEqual(i * i, self.evaluate(get_next())) - ev.wait() - for i in range(event_will_be_set_after_consuming, 100): - self.assertEqual(i * i, self.evaluate(get_next())) - with self.assertRaises(errors.OutOfRangeError): - self.evaluate(get_next()) - - for buffer_size in range(1, set_event_during_invocation): - do_test_ev(buffer_size) - - def testReturnList(self): - dataset = dataset_ops.Dataset.range( - 10).map(lambda x: [x, constant_op.constant(37.0)]) + @combinations.generate(_test_combinations()) + def testReturnList(self, apply_map): + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, lambda x: [x, constant_op.constant(37.0)]) self.assertDatasetProduces( dataset, expected_output=[(i, 37.0) for i in range(10)]) - def testMultiOutputPyFunc(self): + @combinations.generate(_test_combinations()) + def testMultiOutputPyFunc(self, apply_map): # The `tf.py_func()` op returns a list of tensors for its outputs. def _map_fn(x_tensor): def _map_py_func(x): @@ -699,11 +771,13 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): return script_ops.py_func( _map_py_func, [x_tensor], [dtypes.int64, dtypes.float64]) - dataset = dataset_ops.Dataset.range(10).map(_map_fn) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _map_fn) self.assertDatasetProduces( dataset, expected_output=[(i, 37.0) for i in range(10)]) - def testSparse(self): + @combinations.generate(_test_combinations()) + def testSparse(self, apply_map): def _sparse(i): return sparse_tensor.SparseTensorValue( @@ -711,11 +785,13 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): values=(i * np.array([1])), dense_shape=np.array([1, 1])) - dataset = dataset_ops.Dataset.range(10).map(_sparse) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _sparse) self.assertDatasetProduces( dataset, expected_output=[_sparse(i) for i in range(10)]) - def testSparseChain(self): + @combinations.generate(_test_combinations()) + def testSparseChain(self, apply_map): def _sparse(i): return sparse_tensor.SparseTensorValue( @@ -727,37 +803,38 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertTrue(sparse_tensor.is_sparse(i)) return sparse_ops.sparse_concat(0, [i, i]) - dataset = dataset_ops.Dataset.range(10).map(_sparse).map(_check) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _sparse) + dataset = apply_map(dataset, _check) self.assertDatasetProduces( dataset, expected_output=[self.evaluate(_check(_sparse(i))) for i in range(10)]) - def testSparseMapShapeInference(self): - if not context.executing_eagerly(): - self.skipTest("SparseTensor shape inference requires eager mode") + @combinations.generate(_test_combinations_with_mode("eager")) + def testSparseMapShapeInference(self, apply_map): row_lengths = np.random.randint(0, 4, size=128) values = np.ones(np.sum(row_lengths)) sparse = ragged_tensor.RaggedTensor.from_row_lengths( values, row_lengths).to_sparse() dataset = dataset_ops.Dataset.from_tensor_slices(sparse) dataset = dataset.batch(32, drop_remainder=True) - dataset = dataset.map(lambda x: x) + dataset = apply_map(dataset, lambda x: x) self.assertEqual((32, 3), dataset.element_spec.shape) - def testSparseMapShapeInferencePartial(self): - if not context.executing_eagerly(): - self.skipTest("SparseTensor shape inference requires eager mode") + @combinations.generate(_test_combinations_with_mode("eager")) + def testSparseMapShapeInferencePartial(self, apply_map): row_lengths = np.random.randint(0, 4, size=128) values = np.ones(np.sum(row_lengths)) sparse = ragged_tensor.RaggedTensor.from_row_lengths( values, row_lengths).to_sparse() dataset = dataset_ops.Dataset.from_tensor_slices(sparse) dataset = dataset.batch(32, drop_remainder=False) - dataset = dataset.map(lambda x: x) + dataset = apply_map(dataset, lambda x: x) self.assertEqual([None, 3], dataset.element_spec.shape.as_list()) - def testTensorArray(self): + @combinations.generate(_test_combinations()) + def testTensorArray(self, apply_map): def _tensor_array(i): i = math_ops.cast(i, dtypes.int32) @@ -765,11 +842,13 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): tensor_array_ops.TensorArray(dtypes.int32, element_shape=(), size=i) .unstack(math_ops.range(i, dtype=dtypes.int32))) - dataset = dataset_ops.Dataset.range(10).map(_tensor_array) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _tensor_array) self.assertDatasetProduces( dataset, expected_output=[list(range(i)) for i in range(10)]) - def testTensorArrayChain(self): + @combinations.generate(_test_combinations()) + def testTensorArrayChain(self, apply_map): def _tensor_array(i): i = math_ops.cast(i, dtypes.int32) @@ -781,23 +860,28 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertIsInstance(x, tensor_array_ops.TensorArray) return x.identity() - dataset = dataset_ops.Dataset.range(10).map(_tensor_array).map(_check) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _tensor_array) + dataset = apply_map(dataset, _check) self.assertDatasetProduces( dataset, expected_output=[list(range(i)) for i in range(10)]) - def testRagged(self): + @combinations.generate(_test_combinations()) + def testRagged(self, apply_map): def _ragged(i): return ragged_tensor.RaggedTensor.from_tensor(i * [[1]]) - dataset = dataset_ops.Dataset.range(5).map(_ragged) + dataset = dataset_ops.Dataset.range(5) + dataset = apply_map(dataset, _ragged) self.assertDatasetProduces( dataset, expected_output=[ragged_factory_ops.constant([[i]]) for i in range(5)]) - def testRaggedChain(self): + @combinations.generate(_test_combinations()) + def testRaggedChain(self, apply_map): def _ragged(i): return ragged_tensor.RaggedTensor.from_tensor(i * [[1]]) @@ -806,7 +890,9 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertTrue(ragged_tensor.is_ragged(i)) return ragged_concat_ops.concat([i, i], 0) - dataset = dataset_ops.Dataset.range(10).map(_ragged).map(_concat) + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, _ragged) + dataset = apply_map(dataset, _concat) self.assertDatasetProduces( dataset, @@ -815,15 +901,19 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): for i in range(10) ]) - @test_util.run_v1_only("b/123904513") - def testParallelMapOutOfRangeError(self): + # TODO(b/123904513) + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testParallelMapOutOfRangeError(self, apply_map): + def raising_py_func(i): if i == 100: raise StopIteration() else: return i - dataset = dataset_ops.Dataset.range(105).map( + dataset = dataset_ops.Dataset.range(105) + dataset = apply_map( + dataset, lambda x: script_ops.py_func(raising_py_func, [x], dtypes.int64), num_parallel_calls=2) get_next = self.getNext(dataset) @@ -832,11 +922,15 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testConstantOutput(self): - dataset = dataset_ops.Dataset.range(10).map(lambda x: [x, "hello", 10]) + @combinations.generate(_test_combinations()) + def testConstantOutput(self, apply_map): + dataset = dataset_ops.Dataset.range(10) + dataset = apply_map(dataset, lambda x: [x, "hello", 10]) self.assertDatasetProduces(dataset, [(i, b"hello", 10) for i in range(10)]) - def testWarnOnLookupTable(self): + @combinations.generate(_test_combinations()) + def testWarnOnLookupTable(self, apply_map): + def collecting_function(x): _ = lookup_ops.HashTable( lookup_ops.KeyValueTensorInitializer(["a"], [1.]), 0.0, name="t1") @@ -844,30 +938,8 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): warnings.simplefilter("always") with warnings.catch_warnings(record=True) as w: - _ = dataset_ops.Dataset.range(10).map(collecting_function) - # NOTE(mrry): Python 3 prints other warnings in addition to the one we are - # testing, so we search for the expected warning. - self.assertGreaterEqual(len(w), 1) - found_warning = False - for warning in w: - if ("Creating resources inside a function passed to Dataset.map() is " - "not supported." in str(warning)): - found_warning = True - break - self.assertTrue(found_warning) - - @test_util.run_v1_only("map_with_legacy_function v1 only") - def testWarnOnLookupTableLegacyFunction(self): - - def collecting_function(x): - _ = lookup_ops.HashTable( - lookup_ops.KeyValueTensorInitializer(["a"], [1.]), 0.0, name="t1") - return x - - warnings.simplefilter("always") - with warnings.catch_warnings(record=True) as w: - _ = dataset_ops.Dataset.range(10).map_with_legacy_function( - collecting_function) + dataset = dataset_ops.Dataset.range(10) + _ = apply_map(dataset, collecting_function) # NOTE(mrry): Python 3 prints other warnings in addition to the one we are # testing, so we search for the expected warning. self.assertGreaterEqual(len(w), 1) @@ -879,21 +951,25 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): break self.assertTrue(found_warning) + @combinations.generate(test_base.default_test_combinations()) def testWarnOnSeedFromOuterGraph(self): with ops.Graph().as_default() as g: g.seed = 10 warnings.simplefilter("always") + def _check_warning(caught_warnings, expected_result): + found_warning = False + for warning in caught_warnings: + if ("Explicitly set the seed in the function if this is not the " + "intended behavior" in str(warning)): + found_warning = True + break + self.assertEqual(found_warning, expected_result) + # map_fun doesn't use seed, so no warning is generated. with warnings.catch_warnings(record=True) as w: _ = dataset_ops.Dataset.range(10).map(math_ops.square) - found_warning = False - for warning in w: - if ("Explicitly set the seed in the function if this is not the " - "intended behavior" in str(warning)): - found_warning = True - break - self.assertFalse(found_warning) + _check_warning(w, False) def random_func(x): x = math_ops.add(x, 1) @@ -902,14 +978,7 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with warnings.catch_warnings(record=True) as w: _ = dataset_ops.Dataset.range(10).map(random_func) - self.assertGreaterEqual(len(w), 1) - found_warning = False - for warning in w: - if ("Explicitly set the seed in the function if this is not the " - "intended behavior" in str(warning)): - found_warning = True - break - self.assertTrue(found_warning) + _check_warning(w, True) def random_func_seeded(x): ops.get_default_graph().seed = None @@ -918,41 +987,30 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with warnings.catch_warnings(record=True) as w: _ = dataset_ops.Dataset.range(10).batch(2).map(random_func_seeded) - found_warning = False - for warning in w: - if ("Explicitly set the seed in the function if this is not the " - "intended behavior" in str(warning)): - found_warning = True - break - self.assertFalse(found_warning) + _check_warning(w, False) with warnings.catch_warnings(record=True) as w: - _ = dataset_ops.Dataset.range(10).batch( - 2).map(lambda x: random_ops.random_shuffle(x, seed=37)) - found_warning = False - for warning in w: - if ("Explicitly set the seed in the function if this is not the " - "intended behavior" in str(warning)): - found_warning = True - break - self.assertFalse(found_warning) + _ = dataset_ops.Dataset.range(10).batch(2).map( + lambda x: random_ops.random_shuffle(x, seed=37)) + _check_warning(w, False) - def testNestedDatasetMap(self): - # TODO(b/110122868): When iterators can yield a `tf.data.Dataset`, remove - # the `get_single_element()` call. - dataset = dataset_ops.Dataset.from_tensors([1.0, 2.0, 3.0]).map( - dataset_ops.Dataset.from_tensor_slices).map( - lambda ds: ds.batch(3)).flat_map(lambda x: x) + @combinations.generate(_test_combinations()) + def testNestedDatasetMap(self, apply_map): + dataset = dataset_ops.Dataset.from_tensors([1.0, 2.0, 3.0]) + dataset = apply_map(dataset, dataset_ops.Dataset.from_tensor_slices) + dataset = apply_map(dataset, lambda ds: ds.batch(3)).flat_map(lambda x: x) self.assertDatasetProduces(dataset, expected_output=[[1.0, 2.0, 3.0]]) - def testReturnValueError(self): + @combinations.generate(_test_combinations()) + def testReturnValueError(self, apply_map): dataset = dataset_ops.Dataset.from_tensors([1.0, 2.0, 3.0]) with self.assertRaisesRegexp( TypeError, r"Unsupported return value from function passed to " r"Dataset.map\(\): None."): - _ = dataset.map(lambda x: None) + _ = apply_map(dataset, lambda x: None) + @combinations.generate(test_base.default_test_combinations()) def testBrokenFunctionErrorOnInitialization(self): dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0, 3.0]) @@ -965,8 +1023,7 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): value, dtype=dtypes.float32, shape=[0], verify_shape=False)) dtype_value = attr_value_pb2.AttrValue(type=dtypes.int32.as_datatype_enum) - # Create a "Const" op with a `tf.float32` value and a `tf.int32` type - # attr. + # Create a "Const" op with a `tf.float32` value and a `tf.int32` type. const_tensor = ops.get_default_graph().create_op( "Const", [], [dtypes.int32], attrs={ @@ -980,15 +1037,11 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( dataset, expected_error=(errors.InvalidArgumentError, "BrokenConst")) -# pylint: disable=g-long-lambda - @parameterized.named_parameters( - ("Map", lambda dataset, func: - dataset_ops.MapDataset(dataset, func, use_inter_op_parallelism=False)), - ("ParallelMap", lambda dataset, func: - dataset_ops.ParallelMapDataset(dataset, func, num_parallel_calls=1, - use_inter_op_parallelism=False)), - ) - def testNoInterOpParallelism(self, make_dataset_fn): + @combinations.generate( + combinations.times( + _test_combinations_with_mode("graph"), + combinations.combine(num_parallel_calls=[None, 12]))) + def testNoInterOpParallelism(self, apply_map, num_parallel_calls): dataset = dataset_ops.Dataset.from_tensors(0) def _get_tid(): @@ -1000,58 +1053,54 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): tids.append(script_ops.py_func(_get_tid, [], dtypes.int64)) return tids - dataset = make_dataset_fn(dataset, _map_fn) + dataset = apply_map(dataset, _map_fn) + dataset._variant_tensor.op._set_attr("use_inter_op_parallelism", + attr_value_pb2.AttrValue(b=False)) get_next = self.getNext(dataset) tids = self.evaluate(get_next()) self.assertTrue(all(tids[0] == tid for tid in tids)) -# pylint: enable=g-long-lambda - @parameterized.named_parameters( - ("SequentialIdentity", None, lambda x: x, None), - ("SequentialReplicate", None, lambda x: (x, x), None), - ("SequentialSwap", (None, None), lambda x, y: (y, x), None), - ("SequentialProject", (None, None), lambda x, y: x, None), - ("ParallelIdentity", None, lambda x: x, 10), - ("ParallelReplicate", None, lambda x: (x, x), 10), - ("ParallelSwap", (None, None), lambda x, y: (y, x), 10), - ("ParallelProject", (None, None), lambda x, y: x, 10), - ) - def testShortCircuit(self, structure, map_fn, num_parallel_calls): - dataset = self.structuredDataset(structure).repeat().map( - map_fn, num_parallel_calls=num_parallel_calls) + @combinations.generate( + combinations.times(_test_combinations(), _short_circuit_test_cases(), + combinations.combine(num_parallel_calls=[None, 12]))) + def testShortCircuit(self, apply_map, structure, fn, num_parallel_calls): + dataset = self.structuredDataset(structure).repeat() + dataset = apply_map(dataset, fn, num_parallel_calls=num_parallel_calls) get_next = self.getNext(dataset) if isinstance(structure, tuple): - expected = map_fn(*self.evaluate(self.structuredElement(structure))) + expected = fn(*self.evaluate(self.structuredElement(structure))) else: - expected = map_fn(self.evaluate(self.structuredElement(structure))) + expected = fn(self.evaluate(self.structuredElement(structure))) self.assertEqual(expected, self.evaluate(get_next())) - @parameterized.named_parameters( - ("Sequential", None), - ("Parallel", 10), - ) - def testShortCircuitCapturedInput(self, num_parallel_calls): + @combinations.generate( + combinations.times(_test_combinations(), + combinations.combine(num_parallel_calls=[None, 12]))) + def testShortCircuitCapturedInput(self, apply_map, num_parallel_calls): captured_t = variables.Variable(42) - dataset = self.structuredDataset(None).repeat().map( - lambda x: captured_t, num_parallel_calls=num_parallel_calls) + dataset = self.structuredDataset(None).repeat() + dataset = apply_map( + dataset, lambda x: captured_t, num_parallel_calls=num_parallel_calls) self.evaluate(variables.global_variables_initializer()) get_next = self.getNext(dataset, requires_initialization=True) self.assertEqual(42, self.evaluate(get_next())) - @parameterized.named_parameters( - ("1", 1, 1), - ("2", 10, 1), - ("3", 10, 10), - ("4", 100, 1), - ("5", 100, 10), - ("6", 100, 100), - ) - def testSloppyInterleaveInOrder(self, num_elements, num_parallel_calls): + @combinations.generate( + combinations.times( + _test_combinations(), + combinations.combine(num_elements=1, num_parallel_calls=1) + + combinations.combine(num_elements=10, num_parallel_calls=1) + + combinations.combine(num_elements=10, num_parallel_calls=10) + + combinations.combine(num_elements=100, num_parallel_calls=1) + + combinations.combine(num_elements=100, num_parallel_calls=10) + + combinations.combine(num_elements=100, num_parallel_calls=100))) + def testSloppyInterleaveInOrder(self, apply_map, num_elements, + num_parallel_calls): dataset, coordination_events = _make_coordinated_sloppy_dataset( - num_elements, num_parallel_calls) + apply_map, num_elements, num_parallel_calls) options = dataset_ops.Options() options.experimental_threading = threading_options.ThreadingOptions() options.experimental_threading.private_threadpool_size = ( @@ -1064,14 +1113,16 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @parameterized.named_parameters( - ("1", 10, 10), - ("2", 100, 10), - ("3", 100, 100), - ) - def testSloppyInterleaveOutOfOrder(self, num_elements, num_parallel_calls): + @combinations.generate( + combinations.times( + _test_combinations(), + combinations.combine(num_elements=10, num_parallel_calls=10) + + combinations.combine(num_elements=100, num_parallel_calls=10) + + combinations.combine(num_elements=100, num_parallel_calls=100))) + def testSloppyInterleaveOutOfOrder(self, apply_map, num_elements, + num_parallel_calls): dataset, coordination_events = _make_coordinated_sloppy_dataset( - num_elements, num_parallel_calls) + apply_map, num_elements, num_parallel_calls) options = dataset_ops.Options() options.experimental_threading = threading_options.ThreadingOptions() options.experimental_threading.private_threadpool_size = ( @@ -1090,25 +1141,25 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @parameterized.named_parameters( - ("Map", None), - ("ParallelMap", 12), - ) + @combinations.generate( + combinations.combine( + tf_api_version=2, + mode=["eager", "graph"], + num_parallel_calls=[None, 12])) def testPreserveCardinality(self, num_parallel_calls): def py_fn(_): raise StopIteration() - dataset = dataset_ops.DatasetV2.from_tensors(0).map( + dataset = dataset_ops.Dataset.from_tensors(0).map( lambda x: script_ops.py_func(py_fn, [x], dtypes.int64), num_parallel_calls=num_parallel_calls) get_next = self.getNext(dataset) with self.assertRaises(errors.InvalidArgumentError): self.evaluate(get_next()) - # NOTE: collection test is specific to graph mode only, no eager coverage. - @test_util.run_v1_only("graph specific test") - def testSkipEagerCollectionCopy(self): + @combinations.generate(_test_combinations_with_mode("graph")) + def testCollectionCopy(self, apply_map): w = variable_scope.get_variable("w", []) self.assertIn(w, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)) @@ -1117,22 +1168,21 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): return x dataset = dataset_ops.Dataset.from_tensors(constant_op.constant(1.0)) - dataset.map(func) + _ = apply_map(dataset, func) - @parameterized.named_parameters( - ("Sequential", None), - ("Parallel", 12), - ) - @test_util.run_v1_only("graph-mode specific test") - def testSkipEagerMapCancellation(self, num_parallel_calls): + @combinations.generate( + combinations.times( + _test_combinations_with_mode_v1("graph"), + combinations.combine(num_parallel_calls=[None, 12]))) + def testMapCancellation(self, apply_map, num_parallel_calls): # Checks that a cancellation of is threaded through to map transformation. queue = data_flow_ops.FIFOQueue(10, dtypes.int32, ()) def fn(_): return queue.dequeue() - dataset = dataset_ops.Dataset.range(1).map( - fn, num_parallel_calls=num_parallel_calls) + dataset = dataset_ops.Dataset.range(1) + dataset = apply_map(dataset, fn, num_parallel_calls=num_parallel_calls) get_next = self.getNext(dataset, requires_initialization=True) with self.cached_session() as sess: @@ -1143,17 +1193,11 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase): thread.join() -# TODO(shivaniagarwal): separate out `map` and `map_with_legacy_function` tests -# as later would not work in v2. -@test_util.run_all_in_graph_and_eager_modes -class MapWithCapturedVariableTests(test_base.DatasetTestBase, - parameterized.TestCase): - # TODO(b/126553094): map doesnt work with variable defined inside function in # eager mode, possible Graph tensors leak out of the function building context # from function graph in eager mode as variables are created in init_scope. - @test_util.run_v1_only("b/126553094") - def testSkipEagerCreateVariableInsideFunctionWithGetter(self): + @combinations.generate(test_base.graph_only_combinations()) + def testCreateVariableInsideFunctionWithGetter(self): def func(_): with variable_scope.variable_scope( @@ -1162,12 +1206,13 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, "counter", (), dtypes.int32, use_resource=True) return counter_var.assign_add(1) - # NOTE: In the legacy function, resource is captured by value for variable - # getter. dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - with self.assertRaisesWithPredicateMatch( - AttributeError, "'Tensor' object has no attribute 'assign_add'"): - dataset.map_with_legacy_function(func) + + if hasattr(dataset, "map_with_legacy_function"): + # NOTE: In the legacy function, resource is captured by value. + with self.assertRaisesWithPredicateMatch( + AttributeError, "'Tensor' object has no attribute 'assign_add'"): + dataset.map_with_legacy_function(func) dataset = dataset.map(func) self.evaluate(variables.global_variables_initializer()) @@ -1179,18 +1224,12 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @parameterized.named_parameters( - ("MapLegacyFunction", - lambda dataset, func: dataset.map_with_legacy_function(func)), - ("Map", lambda dataset, func: dataset.map(func)), - ) - @test_util.run_v1_only("map_with_legacy_function is only available in v1.") - def testCaptureVariable(self, transformation_function): + @combinations.generate(_test_combinations()) + def testCaptureVariable(self, apply_map): counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = transformation_function( - dataset, lambda _: counter_var.assign_add(1)) + dataset = apply_map(dataset, lambda _: counter_var.assign_add(1)) get_next = self.getNext(dataset, requires_initialization=True) self.evaluate(counter_var.initializer) @@ -1203,34 +1242,20 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, self.evaluate(get_next()) self.assertEqual(10, self.evaluate(counter_var)) - # NOTE: no need to explicitly initialize variables in eager mode. - @parameterized.named_parameters( - ("MapLegacyFunction", - lambda dataset, func: dataset.map_with_legacy_function(func)), - ("Map", lambda dataset, func: dataset.map(func)), - ) - @test_util.run_v1_only("this test is meant to run in graph mode only.") - def testSkipEagerCaptureUninitializedVariableError(self, - transformation_function): + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testCaptureUninitializedVariableError(self, apply_map): counter_var = variable_scope.get_variable( "counter", (), dtypes.int32, use_resource=True) dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = transformation_function( - dataset, lambda _: counter_var.assign_add(1)) + dataset = apply_map(dataset, lambda _: counter_var.assign_add(1)) get_next = self.getNext(dataset, requires_initialization=True) with self.assertRaises(errors.NotFoundError): self.evaluate(get_next()) # TODO(b/121264236): add eager mode coverage when we have multi-device setup. - @parameterized.named_parameters( - ("MapLegacyFunction", - lambda dataset, func: dataset.map_with_legacy_function(func)), - ("Map", lambda dataset, func: dataset.map(func)), - ) - @test_util.run_v1_only("b/121264236") - def testSkipEagerCaptureConstantsWithConflictingDevices( - self, transformation_function): + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testCaptureConstantsWithConflictingDevices(self, apply_map): config = config_pb2.ConfigProto(device_count={"CPU": 3}) with self.cached_session(config=config): with ops.device("/device:CPU:0"): @@ -1242,13 +1267,13 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, return math_ops.add(a, b) dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = transformation_function(dataset, func) + dataset = apply_map(dataset, func) expected_output = [8.0] * 10 self.assertDatasetProduces(dataset, expected_output=expected_output) # TODO(b/121264236): add eager mode coverage when we have multi-device setup. - @test_util.run_v1_only("b/121264236") - def testSkipEagerRefVariablesWithMultipleDevices(self): + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testReferenceVariablesWithMultipleDevices(self, apply_map): config = config_pb2.ConfigProto(device_count={"CPU": 3}) with self.cached_session(config=config): @@ -1262,7 +1287,7 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, # NOTE: Use the legacy function implementation as eager function will # convert RefVariables to ResourceVariables. dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = dataset.map_with_legacy_function(func) + dataset = apply_map(dataset, func) self.evaluate(variables.global_variables_initializer()) expected_output = [8.0] * 10 self.assertDatasetProduces( @@ -1271,8 +1296,8 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, requires_initialization=True) # TODO(b/121264236): add eager mode coverage when we have multi-device setup. - @test_util.run_v1_only("b/121264236") - def testSkipEagerResourceVariablesWithMultipleDevices(self): + @combinations.generate(_test_combinations_with_mode_v1("graph")) + def testResourceVariablesWithMultipleDevices(self, apply_map): config = config_pb2.ConfigProto(device_count={"CPU": 3}) def func(_): @@ -1287,25 +1312,10 @@ class MapWithCapturedVariableTests(test_base.DatasetTestBase, "b", (), dtypes.int32, use_resource=True) return math_ops.add(a_var, b_var) - g_1 = ops.Graph() - with self.session(config=config, graph=g_1): - # The MapDataset node ends up with two ResourceVariable inputs, one on - # device CPU:0 and the other on device CPU:1. + g = ops.Graph() + with self.session(config=config, graph=g): dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = dataset.map(func) - self.evaluate(variables.global_variables_initializer()) - expected_output = [1] * 10 - self.assertDatasetProduces( - dataset, - expected_output=expected_output, - requires_initialization=True) - - g_2 = ops.Graph() - with self.session(config=config, graph=g_2): - # In old-Defun variable is captured as value, hence there is no colocation - # error. - dataset = dataset_ops.Dataset.from_tensors(0).repeat(10) - dataset = dataset.map_with_legacy_function(func) + dataset = apply_map(dataset, func) self.evaluate(variables.global_variables_initializer()) expected_output = [1] * 10 self.assertDatasetProduces( diff --git a/tensorflow/python/data/kernel_tests/memory_cleanup_test.py b/tensorflow/python/data/kernel_tests/memory_cleanup_test.py index a2015ef47d1..5b0ea02a054 100644 --- a/tensorflow/python/data/kernel_tests/memory_cleanup_test.py +++ b/tensorflow/python/data/kernel_tests/memory_cleanup_test.py @@ -119,8 +119,7 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase): ] self.assertEmpty(tensors, "%d Tensors are still alive." % len(tensors)) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testFilter(self): def get_dataset(): @@ -144,8 +143,7 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase): self._testIteratorMemoryLeak(get_dataset) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testFlatMap(self): def get_dataset(): @@ -157,8 +155,7 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase): self._testIteratorMemoryLeak(get_dataset) - @combinations.generate( - combinations.combine(tf_api_version=[1, 2], mode="eager")) + @combinations.generate(test_base.eager_only_combinations()) def testFromGenerator(self): def get_dataset(): @@ -171,8 +168,8 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase): self._testIteratorMemoryLeak(get_dataset) @combinations.generate( - combinations.combine( - tf_api_version=[1, 2], mode="eager", num_parallel_calls=[None, 10])) + combinations.times(test_base.eager_only_combinations(), + combinations.combine(num_parallel_calls=[None, 10]))) def testMap(self, num_parallel_calls): def get_dataset(): @@ -201,8 +198,8 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase): self._testIteratorMemoryLeak(get_dataset) @combinations.generate( - combinations.combine( - tf_api_version=[1, 2], mode="eager", num_parallel_calls=[None, 10])) + combinations.times(test_base.eager_only_combinations(), + combinations.combine(num_parallel_calls=[None, 10]))) def testInterleave(self, num_parallel_calls): def get_dataset(): diff --git a/tensorflow/python/data/kernel_tests/optional_test.py b/tensorflow/python/data/kernel_tests/optional_test.py index 3ab6717b9c3..f0795563d09 100644 --- a/tensorflow/python/data/kernel_tests/optional_test.py +++ b/tensorflow/python/data/kernel_tests/optional_test.py @@ -17,6 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools + from absl.testing import parameterized import numpy as np @@ -27,6 +29,7 @@ from tensorflow.python.data.ops import optional_ops from tensorflow.python.data.util import structure from tensorflow.python.eager import context from tensorflow.python.eager import def_function +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -40,14 +43,90 @@ from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes +def _optional_spec_test_combinations(): + # pylint: disable=g-long-lambda + cases = [ + ("Dense", lambda: constant_op.constant(37.0), + tensor_spec.TensorSpec([], dtypes.float32)), + ("Sparse", lambda: sparse_tensor.SparseTensor( + indices=[[0, 1]], + values=constant_op.constant([0], dtype=dtypes.int32), + dense_shape=[10, 10]), + sparse_tensor.SparseTensorSpec([10, 10], dtypes.int32)), + ("Nest", lambda: { + "a": constant_op.constant(37.0), + "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar")) + }, { + "a": + tensor_spec.TensorSpec([], dtypes.float32), + "b": ( + tensor_spec.TensorSpec([1], dtypes.string), + tensor_spec.TensorSpec([], dtypes.string), + ) + }), + ("Optional", lambda: optional_ops.Optional.from_value(37.0), + optional_ops.OptionalSpec(tensor_spec.TensorSpec([], dtypes.float32))), + ] + + def reduce_fn(x, y): + name, value_fn, expected_structure = y + return x + combinations.combine( + tf_value_fn=combinations.NamedObject(name, value_fn), + expected_value_structure=expected_structure) + + return functools.reduce(reduce_fn, cases, []) + + +def _get_next_as_optional_test_combinations(): + # pylint: disable=g-long-lambda + cases = [ + ("Dense", np.array([1, 2, 3], dtype=np.int32), + lambda: constant_op.constant([4, 5, 6], dtype=dtypes.int32), True), + ("Sparse", + sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 1]], + values=np.array([-1., 1.], dtype=np.float32), + dense_shape=[2, 2]), + lambda: sparse_tensor.SparseTensor( + indices=[[0, 1], [1, 0]], values=[37.0, 42.0], dense_shape=[2, 2]), + False), + ("Nest", { + "a": + np.array([1, 2, 3], dtype=np.int32), + "b": + sparse_tensor.SparseTensorValue( + indices=[[0, 0], [1, 1]], + values=np.array([-1., 1.], dtype=np.float32), + dense_shape=[2, 2]) + }, lambda: { + "a": + constant_op.constant([4, 5, 6], dtype=dtypes.int32), + "b": + sparse_tensor.SparseTensor( + indices=[[0, 1], [1, 0]], + values=[37.0, 42.0], + dense_shape=[2, 2]) + }, False), + ] + + def reduce_fn(x, y): + name, value, value_fn, gpu_compatible = y + return x + combinations.combine( + np_value=value, tf_value_fn=combinations.NamedObject(name, value_fn), + gpu_compatible=gpu_compatible) + + return functools.reduce(reduce_fn, cases, []) + + class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): + @combinations.generate(test_base.default_test_combinations()) def testFromValue(self): opt = optional_ops.Optional.from_value(constant_op.constant(37.0)) self.assertTrue(self.evaluate(opt.has_value())) self.assertEqual(37.0, self.evaluate(opt.get_value())) + @combinations.generate(test_base.default_test_combinations()) def testFromStructuredValue(self): opt = optional_ops.Optional.from_value({ "a": constant_op.constant(37.0), @@ -59,6 +138,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): "b": ([b"Foo"], b"Bar") }, self.evaluate(opt.get_value())) + @combinations.generate(test_base.default_test_combinations()) def testFromSparseTensor(self): st_0 = sparse_tensor.SparseTensorValue( indices=np.array([[0]]), @@ -77,6 +157,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertAllEqual(expected.dense_shape, self.evaluate(actual.dense_shape)) + @combinations.generate(test_base.default_test_combinations()) def testFromNone(self): value_structure = tensor_spec.TensorSpec([], dtypes.float32) opt = optional_ops.Optional.none_from_structure(value_structure) @@ -91,6 +172,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(opt.get_value()) + @combinations.generate(test_base.default_test_combinations()) def testAddN(self): devices = ["/cpu:0"] if test_util.is_gpu_available(): @@ -117,6 +199,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): opt_none1.value_structure) self.assertFalse(self.evaluate(add_opt.has_value())) + @combinations.generate(test_base.default_test_combinations()) def testNestedAddN(self): devices = ["/cpu:0"] if test_util.is_gpu_available(): @@ -137,6 +220,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): opt1.value_structure) self.assertAllEqual(inner_add_opt.get_value(), [4, 6.0]) + @combinations.generate(test_base.default_test_combinations()) def testZerosLike(self): devices = ["/cpu:0"] if test_util.is_gpu_available(): @@ -159,6 +243,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): opt_none.value_structure) self.assertFalse(self.evaluate(zeros_opt.has_value())) + @combinations.generate(test_base.default_test_combinations()) def testNestedZerosLike(self): devices = ["/cpu:0"] if test_util.is_gpu_available(): @@ -175,6 +260,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): opt1.value_structure) self.assertEqual(self.evaluate(inner_zeros_opt.get_value()), 0.0) + @combinations.generate(test_base.default_test_combinations()) def testCopyToGPU(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") @@ -204,6 +290,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): self.evaluate(gpu_optional_with_value_values)) self.assertFalse(self.evaluate(gpu_optional_none_has_value)) + @combinations.generate(test_base.default_test_combinations()) def testNestedCopyToGPU(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") @@ -239,42 +326,10 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertFalse(self.evaluate(inner_none.has_value())) self.assertEqual(1.0, self.evaluate(gpu_nested_optional_values[2])) - def _assertElementValueEqual(self, expected, actual): - if isinstance(expected, dict): - self.assertItemsEqual(list(expected.keys()), list(actual.keys())) - for k in expected.keys(): - self._assertElementValueEqual(expected[k], actual[k]) - elif isinstance(expected, sparse_tensor.SparseTensorValue): - self.assertAllEqual(expected.indices, actual.indices) - self.assertAllEqual(expected.values, actual.values) - self.assertAllEqual(expected.dense_shape, actual.dense_shape) - else: - self.assertAllEqual(expected, actual) - - # pylint: disable=g-long-lambda - @parameterized.named_parameters( - ("Tensor", lambda: constant_op.constant(37.0), - tensor_spec.TensorSpec([], dtypes.float32)), - ("SparseTensor", lambda: sparse_tensor.SparseTensor( - indices=[[0, 1]], - values=constant_op.constant([0], dtype=dtypes.int32), - dense_shape=[10, 10]), - sparse_tensor.SparseTensorSpec([10, 10], dtypes.int32)), - ("Nest", lambda: { - "a": constant_op.constant(37.0), - "b": (constant_op.constant(["Foo"]), constant_op.constant("Bar")) - }, { - "a": - tensor_spec.TensorSpec([], dtypes.float32), - "b": ( - tensor_spec.TensorSpec([1], dtypes.string), - tensor_spec.TensorSpec([], dtypes.string), - ) - }), - ("Optional", lambda: optional_ops.Optional.from_value(37.0), - optional_ops.OptionalSpec( - tensor_spec.TensorSpec([], dtypes.float32))), - ) + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + _optional_spec_test_combinations())) def testOptionalSpec(self, tf_value_fn, expected_value_structure): tf_value = tf_value_fn() opt = optional_ops.Optional.from_value(tf_value) @@ -304,36 +359,21 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): round_trip_opt = opt_structure._from_tensor_list( opt_structure._to_tensor_list(opt)) if isinstance(tf_value, optional_ops.Optional): - self._assertElementValueEqual( + self.assertValuesEqual( self.evaluate(tf_value.get_value()), self.evaluate(round_trip_opt.get_value().get_value())) else: - self._assertElementValueEqual( + self.assertValuesEqual( self.evaluate(tf_value), self.evaluate(round_trip_opt.get_value())) - @parameterized.named_parameters( - ("Tensor", np.array([1, 2, 3], dtype=np.int32), - lambda: constant_op.constant([4, 5, 6], dtype=dtypes.int32), True), - ("SparseTensor", sparse_tensor.SparseTensorValue( - indices=[[0, 0], [1, 1]], - values=np.array([-1., 1.], dtype=np.float32), dense_shape=[2, 2]), - lambda: sparse_tensor.SparseTensor( - indices=[[0, 1], [1, 0]], values=[37.0, 42.0], dense_shape=[2, 2]), - False), - ("Nest", {"a": np.array([1, 2, 3], dtype=np.int32), - "b": sparse_tensor.SparseTensorValue( - indices=[[0, 0], [1, 1]], - values=np.array([-1., 1.], dtype=np.float32), - dense_shape=[2, 2])}, - lambda: {"a": constant_op.constant([4, 5, 6], dtype=dtypes.int32), - "b": sparse_tensor.SparseTensor( - indices=[[0, 1], [1, 0]], values=[37.0, 42.0], - dense_shape=[2, 2])}, False), - ) + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + _get_next_as_optional_test_combinations())) def testIteratorGetNextAsOptional(self, np_value, tf_value_fn, - works_on_gpu): - if not works_on_gpu and test.is_gpu_available(): + gpu_compatible): + if not gpu_compatible and test.is_gpu_available(): self.skipTest("Test case not yet supported on GPU.") ds = dataset_ops.Dataset.from_tensors(np_value).repeat(3) @@ -348,7 +388,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): next_elem.value_structure, structure.type_spec_from_value(tf_value_fn()))) self.assertTrue(next_elem.has_value()) - self._assertElementValueEqual(np_value, next_elem.get_value()) + self.assertValuesEqual(np_value, next_elem.get_value()) # After exhausting the iterator, `next_elem.has_value()` will evaluate to # false, and attempting to get the value will fail. for _ in range(2): @@ -379,7 +419,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): elem_has_value, elem_value = self.evaluate( [elem_has_value_t, elem_value_t]) self.assertTrue(elem_has_value) - self._assertElementValueEqual(np_value, elem_value) + self.assertValuesEqual(np_value, elem_value) # After exhausting the iterator, `next_elem.has_value()` will evaluate to # false, and attempting to get the value will fail. @@ -388,6 +428,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(elem_value_t) + @combinations.generate(test_base.default_test_combinations()) def testFunctionBoundaries(self): @def_function.function def get_optional(): @@ -407,6 +448,7 @@ class OptionalTest(test_base.DatasetTestBase, parameterized.TestCase): val = consume_optional(opt_tensor) self.assertEqual(self.evaluate(val), 1.0) + @combinations.generate(test_base.default_test_combinations()) def testLimitedRetracing(self): trace_count = [0] diff --git a/tensorflow/python/data/kernel_tests/options_test.py b/tensorflow/python/data/kernel_tests/options_test.py index 222d8c6f1a6..b38d008b833 100644 --- a/tensorflow/python/data/kernel_tests/options_test.py +++ b/tensorflow/python/data/kernel_tests/options_test.py @@ -18,25 +18,31 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized + from tensorflow.python.data.experimental.ops import optimization_options from tensorflow.python.data.experimental.ops import stats_options from tensorflow.python.data.experimental.ops import threading_options from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.platform import test -class OptionsTest(test_base.DatasetTestBase): +class OptionsTest(test_base.DatasetTestBase, parameterized.TestCase): + @combinations.generate(test_base.default_test_combinations()) def testOptionsDefault(self): ds = dataset_ops.Dataset.range(0) self.assertEqual(dataset_ops.Options(), ds.options()) + @combinations.generate(test_base.default_test_combinations()) def testOptionsOnce(self): options = dataset_ops.Options() ds = dataset_ops.Dataset.range(0).with_options(options).cache() self.assertEqual(options, ds.options()) + @combinations.generate(test_base.default_test_combinations()) def testOptionsTwiceSame(self): options = dataset_ops.Options() options.experimental_optimization.autotune = True @@ -44,6 +50,7 @@ class OptionsTest(test_base.DatasetTestBase): options) self.assertEqual(options, ds.options()) + @combinations.generate(test_base.default_test_combinations()) def testOptionsTwiceDifferent(self): options1 = dataset_ops.Options() options1.experimental_optimization.autotune = True @@ -55,6 +62,7 @@ class OptionsTest(test_base.DatasetTestBase): # Explicitly check that flag is False since assertFalse allows None self.assertIs(ds.options().experimental_deterministic, False) + @combinations.generate(test_base.default_test_combinations()) def testOptionsTwiceDifferentError(self): options1 = dataset_ops.Options() options1.experimental_optimization.autotune = True @@ -64,6 +72,7 @@ class OptionsTest(test_base.DatasetTestBase): "Cannot merge incompatible values"): dataset_ops.Dataset.range(0).with_options(options1).with_options(options2) + @combinations.generate(test_base.default_test_combinations()) def testOptionsMergeOptionsFromMultipleInputs(self): options1 = dataset_ops.Options() options1.experimental_optimization.autotune = True @@ -75,6 +84,7 @@ class OptionsTest(test_base.DatasetTestBase): self.assertTrue(ds.options().experimental_optimization.autotune) self.assertTrue(ds.options().experimental_deterministic) + @combinations.generate(test_base.default_test_combinations()) def testOptionsHaveDefaults(self): options1 = dataset_ops.Options() options2 = dataset_ops.Options() @@ -84,12 +94,11 @@ class OptionsTest(test_base.DatasetTestBase): options2.experimental_stats) self.assertIsNot(options1.experimental_threading, options2.experimental_threading) - self.assertEquals(options1.experimental_optimization, - optimization_options.OptimizationOptions()) - self.assertEquals(options1.experimental_stats, - stats_options.StatsOptions()) - self.assertEquals(options1.experimental_threading, - threading_options.ThreadingOptions()) + self.assertEqual(options1.experimental_optimization, + optimization_options.OptimizationOptions()) + self.assertEqual(options1.experimental_stats, stats_options.StatsOptions()) + self.assertEqual(options1.experimental_threading, + threading_options.ThreadingOptions()) if __name__ == "__main__": diff --git a/tensorflow/python/data/kernel_tests/padded_batch_test.py b/tensorflow/python/data/kernel_tests/padded_batch_test.py index 39339c0063a..a3b8f3945f3 100644 --- a/tensorflow/python/data/kernel_tests/padded_batch_test.py +++ b/tensorflow/python/data/kernel_tests/padded_batch_test.py @@ -23,43 +23,30 @@ import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_shape -from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import string_ops from tensorflow.python.platform import test from tensorflow.python.util import compat -def _random_seq_lens(count): - return np.random.randint(20, size=(count,)).astype(np.int32) - - -@test_util.run_all_in_graph_and_eager_modes class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): - @parameterized.named_parameters( - ('default_padding', _random_seq_lens(32), 4, [-1], False), - ('constant_padding', _random_seq_lens(32), 4, [25], False), - ('uneven_with_remainder', _random_seq_lens(34), 4, [-1], False), - ('uneven_without_remainder', _random_seq_lens(34), 4, [-1], True), - ) - def testPaddedBatchDataset(self, seq_lens, batch_size, padded_shapes, - drop_remainder): - """Tests the padded batch dataset logic for various input configurations. - - Args: - seq_lens: the input sequence lengths - batch_size: the batch size - padded_shapes: the padded shapes to use - drop_remainder: whether a smaller batch size should be produced if batch - size does not divide number of inputs evenly - """ - + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + count=[32, 34], + padded_shapes=[[None], [25]], + drop_remainder=[True, False]))) + def testPaddedBatchDataset(self, count, padded_shapes, drop_remainder): + seq_lens = np.random.randint(20, size=(count,)).astype(np.int32) + batch_size = 4 dataset = dataset_ops.Dataset.from_tensor_slices(seq_lens).map( lambda x: array_ops.fill([x], x)).padded_batch( batch_size=batch_size, @@ -81,7 +68,9 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): if not drop_remainder and len(seq_lens) % batch_size > 0: result = self.evaluate(get_next()) - padded_len = np.max(result) if result.size > 0 else 0 + padded_len = padded_shapes[0] + if padded_len is None or padded_len == -1: + padded_len = np.max(result) if result.size > 0 else 0 self.assertEqual((len(seq_lens) % batch_size, padded_len), result.shape) for j in range(len(seq_lens) % batch_size): seq_len = seq_lens[num_full_batches * batch_size + j] @@ -93,7 +82,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @test_util.run_deprecated_v1 + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchShortPadding(self): dataset = ( dataset_ops.Dataset.from_tensor_slices( @@ -102,6 +91,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( dataset, expected_error=(errors.DataLossError, '')) + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchEmptyTensors(self): dataset = ( dataset_ops.Dataset.from_tensor_slices( @@ -109,6 +99,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): batch_size=4, padded_shapes=[-1])) self.assertDatasetProduces(dataset, expected_output=[[[], [], [], []]]) + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchDatasetNonDefaultPadding(self): def fill_tuple(x): @@ -139,6 +130,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchDatasetUnicode(self): # See GitHub issue 16149 def generator(): @@ -156,9 +148,8 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): next_element = self.getNext(padded_dataset) self.evaluate(next_element()) - # NOTE: This test is specific to graph mode and is skipped in eager mode. - @test_util.run_deprecated_v1 - def testSkipEagerPaddedBatchDatasetShapeSpecifications(self): + @combinations.generate(test_base.graph_only_combinations()) + def testPaddedBatchDatasetShapeSpecifications(self): int_placeholder = array_ops.placeholder(dtypes.int32) float_placeholder = array_ops.placeholder(dtypes.float32) string_placeholder = array_ops.placeholder(dtypes.string) @@ -190,6 +181,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertEqual([None, None, None], dataset_output_shapes[1].as_list()) self.assertEqual([None, 37], dataset_output_shapes[2].as_list()) + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchSparseError(self): def _map_fn(i): @@ -199,6 +191,7 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(TypeError): _ = dataset_ops.Dataset.range(10).map(_map_fn).padded_batch(10) + @combinations.generate(test_base.default_test_combinations()) def testPaddedBatchShapeError(self): with self.assertRaisesRegexp( ValueError, r'The padded shape \(1,\) is not compatible with the ' @@ -230,9 +223,8 @@ class PaddedBatchTest(test_base.DatasetTestBase, parameterized.TestCase): _ = dataset_ops.Dataset.range(10).padded_batch( 5, padded_shapes=shape_as_tensor) - # NOTE: This test is specific to graph mode and is skipped in eager mode. - @test_util.run_deprecated_v1 - def testSkipEagerPaddedBatchShapeError(self): + @combinations.generate(test_base.graph_only_combinations()) + def testPaddedBatchShapeErrorPlaceholder(self): with self.assertRaisesRegexp( ValueError, r'The padded shape \((\?|None), (\?|None)\) is not compatible with the ' diff --git a/tensorflow/python/data/kernel_tests/prefetch_test.py b/tensorflow/python/data/kernel_tests/prefetch_test.py index 427fbf1d29f..c6d2877ee7c 100644 --- a/tensorflow/python/data/kernel_tests/prefetch_test.py +++ b/tensorflow/python/data/kernel_tests/prefetch_test.py @@ -23,36 +23,41 @@ from absl.testing import parameterized from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import errors -from tensorflow.python.framework import test_util from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes class PrefetchTest(test_base.DatasetTestBase, parameterized.TestCase): - @parameterized.parameters((-1), (0), (5)) + @combinations.generate( + combinations.times(test_base.default_test_combinations(), + combinations.combine(buffer_size=[-1, None, 0, 42]))) def testBufferSize(self, buffer_size): dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size) self.assertDatasetProduces(dataset, expected_output=range(10)) - @parameterized.parameters((-2), (-42)) + @combinations.generate( + combinations.times(test_base.default_test_combinations(), + combinations.combine(buffer_size=[-2, -42]))) def testInvalidBufferSize(self, buffer_size): with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(10).prefetch(buffer_size=buffer_size) self.evaluate(dataset._variant_tensor) - @parameterized.parameters(*[(buffer_size, slack_period) - for buffer_size in (-1, None, 0, 5) - for slack_period in (1, 8)]) + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + buffer_size=[-1, None, 0, 42], slack_period=[1, 8]))) def testPrefetchWithSlack(self, buffer_size, slack_period): dataset = dataset_ops.Dataset.range(100) dataset = dataset_ops.PrefetchDataset( dataset, buffer_size, slack_period=slack_period) self.assertDatasetProduces(dataset, expected_output=range(100)) - @test_util.run_v1_only("graph-mode specific test") - def testSkipEagerPrefetchCancellation(self): + @combinations.generate(combinations.combine(tf_api_version=1, mode="graph")) + def testPrefetchCancellation(self): def map_py_fn(x): while x > -1: diff --git a/tensorflow/python/data/kernel_tests/range_test.py b/tensorflow/python/data/kernel_tests/range_test.py index b7ac60c3fff..d136565ce42 100644 --- a/tensorflow/python/data/kernel_tests/range_test.py +++ b/tensorflow/python/data/kernel_tests/range_test.py @@ -17,51 +17,60 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized + from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import errors -from tensorflow.python.framework import test_util from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes -class RangeTest(test_base.DatasetTestBase): +class RangeTest(test_base.DatasetTestBase, parameterized.TestCase): + @combinations.generate(test_base.default_test_combinations()) def testStop(self): dataset = dataset_ops.Dataset.range(5) self.assertDatasetProduces(dataset, expected_output=range(5)) + @combinations.generate(test_base.default_test_combinations()) def testStartStop(self): start, stop = 2, 5 dataset = dataset_ops.Dataset.range(start, stop) self.assertDatasetProduces(dataset, expected_output=range(2, 5)) + @combinations.generate(test_base.default_test_combinations()) def testStartStopStep(self): start, stop, step = 2, 10, 2 dataset = dataset_ops.Dataset.range(start, stop, step) self.assertDatasetProduces(dataset, expected_output=range(2, 10, 2)) + @combinations.generate(test_base.default_test_combinations()) def testZeroStep(self): start, stop, step = 2, 10, 0 with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(start, stop, step) self.evaluate(dataset._variant_tensor) + @combinations.generate(test_base.default_test_combinations()) def testNegativeStep(self): start, stop, step = 2, 10, -1 dataset = dataset_ops.Dataset.range(start, stop, step) self.assertDatasetProduces(dataset, expected_output=range(2, 10, -1)) + @combinations.generate(test_base.default_test_combinations()) def testStopLessThanStart(self): start, stop = 10, 2 dataset = dataset_ops.Dataset.range(start, stop) self.assertDatasetProduces(dataset, expected_output=range(10, 2)) + @combinations.generate(test_base.default_test_combinations()) def testStopLessThanStartWithPositiveStep(self): start, stop, step = 10, 2, 2 dataset = dataset_ops.Dataset.range(start, stop, step) self.assertDatasetProduces(dataset, expected_output=range(10, 2, 2)) + @combinations.generate(test_base.default_test_combinations()) def testStopLessThanStartWithNegativeStep(self): start, stop, step = 10, 2, -1 dataset = dataset_ops.Dataset.range(start, stop, step) diff --git a/tensorflow/python/data/kernel_tests/repeat_test.py b/tensorflow/python/data/kernel_tests/repeat_test.py index 8a8537b30cf..c4262fcc08c 100644 --- a/tensorflow/python/data/kernel_tests/repeat_test.py +++ b/tensorflow/python/data/kernel_tests/repeat_test.py @@ -17,43 +17,33 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.framework import test_util +from tensorflow.python.framework import combinations from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes -class RepeatTest(test_base.DatasetTestBase): +class RepeatTest(test_base.DatasetTestBase, parameterized.TestCase): - def testRepeatTensorDataset(self): + @combinations.generate( + combinations.times(test_base.default_test_combinations(), + combinations.combine(count=[0, 3, 7]))) + def testFiniteRepeat(self, count): """Test a dataset that repeats its input multiple times.""" components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) - # This placeholder can be fed when dataset-definition subgraph - # runs (i.e. `init_op` below) to configure the number of - # repetitions used in a particular iterator. + dataset = dataset_ops.Dataset.from_tensors(components).repeat(count) + self.assertEqual( + [c.shape for c in components], + [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) + self.assertDatasetProduces(dataset, [components] * count) - def do_test(count): - dataset = dataset_ops.Dataset.from_tensors(components).repeat(count) - self.assertEqual( - [c.shape for c in components], - [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) - self.assertDatasetProduces(dataset, [components] * count) - - # Test a finite repetition. - do_test(3) - - # test a different finite repetition. - do_test(7) - - # Test an empty repetition. - do_test(0) - - # Test an infinite repetition. - # NOTE(mrry): There's not a good way to test that the sequence - # actually is infinite. + @combinations.generate(test_base.default_test_combinations()) + def testInfiniteRepeat(self): + # NOTE(mrry): There's not a good way to test that the sequence is infinite. + components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) dataset = dataset_ops.Dataset.from_tensors(components).repeat(-1) self.assertEqual( [c.shape for c in components], @@ -64,7 +54,8 @@ class RepeatTest(test_base.DatasetTestBase): for component, result_component in zip(components, results): self.assertAllEqual(component, result_component) - def testRepeatRepeatTensorDataset(self): + @combinations.generate(test_base.default_test_combinations()) + def testRepeatRepeat(self): """Test the composition of repeat datasets.""" components = (np.array(1), np.array([1, 2, 3]), np.array(37.0)) inner_count, outer_count = 7, 14 @@ -77,11 +68,6 @@ class RepeatTest(test_base.DatasetTestBase): self.assertDatasetProduces(dataset, [components] * (inner_count * outer_count)) - def testRepeatEmptyDataset(self): - """Test that repeating an empty dataset does not hang.""" - dataset = dataset_ops.Dataset.from_tensors(0).repeat(10).skip(10).repeat(-1) - self.assertDatasetProduces(dataset, []) - if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/shard_test.py b/tensorflow/python/data/kernel_tests/shard_test.py index 9fc70ff6075..5830b66d61c 100644 --- a/tensorflow/python/data/kernel_tests/shard_test.py +++ b/tensorflow/python/data/kernel_tests/shard_test.py @@ -17,66 +17,79 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized + from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import errors -from tensorflow.python.framework import test_util from tensorflow.python.platform import test -@test_util.run_v1_only("deprecated API, no eager or V2 test coverage") -class ShardTest(test_base.DatasetTestBase): +class ShardTest(test_base.DatasetTestBase, parameterized.TestCase): + @combinations.generate(test_base.default_test_combinations()) def testSimpleCase(self): dataset = dataset_ops.Dataset.range(10).shard(5, 2) self.assertDatasetProduces(dataset, expected_output=[2, 7]) + @combinations.generate(test_base.default_test_combinations()) def testNestedData(self): dataset_a = dataset_ops.Dataset.range(10) dataset_b = dataset_ops.Dataset.range(10, 0, -1) dataset = dataset_ops.Dataset.zip((dataset_a, dataset_b)).shard(5, 2) self.assertDatasetProduces(dataset, expected_output=[(2, 8), (7, 3)]) + @combinations.generate(test_base.default_test_combinations()) def testOffsetZero(self): dataset = dataset_ops.Dataset.range(10).shard(5, 0) self.assertDatasetProduces(dataset, expected_output=[0, 5]) + @combinations.generate(test_base.default_test_combinations()) def testOffsetGreaterNumShards(self): with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(10).shard(5, 7) self.evaluate(self.getNext(dataset)()) + @combinations.generate(test_base.default_test_combinations()) def testNegativeOffset(self): with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(10).shard(5, -3) self.evaluate(self.getNext(dataset)()) + @combinations.generate(test_base.default_test_combinations()) def testNegativeNumShards(self): with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(10).shard(-3, 1) self.evaluate(self.getNext(dataset)()) + @combinations.generate(test_base.default_test_combinations()) def testZeroNumShards(self): with self.assertRaises(errors.InvalidArgumentError): dataset = dataset_ops.Dataset.range(10).shard(0, 1) self.evaluate(self.getNext(dataset)()) + @combinations.generate(test_base.default_test_combinations()) def testIteratorEndsBeforeFirstElem(self): dataset = dataset_ops.Dataset.range(1).shard(5, 2) self.assertDatasetProduces(dataset, expected_output=[]) + @combinations.generate(test_base.default_test_combinations()) def testLargerWorkerPool(self): dataset = dataset_ops.Dataset.range(10).shard(7, 5) self.assertDatasetProduces(dataset, expected_output=[5]) + @combinations.generate(test_base.default_test_combinations()) def testIndexEqualsNumShards(self): dataset = dataset_ops.Dataset.range(10).shard(5, 4) self.assertDatasetProduces(dataset, expected_output=[4, 9]) + @combinations.generate(test_base.default_test_combinations()) def testIndexEqualsNumShards2(self): dataset = dataset_ops.Dataset.range(10).shard(4, 3) self.assertDatasetProduces(dataset, expected_output=[3, 7]) + @combinations.generate(test_base.default_test_combinations()) def testNumShardsLargerThanDataset(self): dataset = dataset_ops.Dataset.range(10).shard(20, 5) self.assertDatasetProduces(dataset, expected_output=[5]) diff --git a/tensorflow/python/data/kernel_tests/shuffle_test.py b/tensorflow/python/data/kernel_tests/shuffle_test.py index 7f801e1b5f4..c9d17b79016 100644 --- a/tensorflow/python/data/kernel_tests/shuffle_test.py +++ b/tensorflow/python/data/kernel_tests/shuffle_test.py @@ -40,7 +40,7 @@ from tensorflow.python.platform import test class ShuffleTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate(test_base.default_test_combinations()) - def testShuffleDataset(self): + def testBasic(self): components = ( np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]), np.array([9.0, 10.0, 11.0, 12.0]) @@ -160,7 +160,7 @@ class ShuffleTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate( combinations.times( - combinations.combine(tf_api_version=[1, 2], mode="graph"), + test_base.graph_only_combinations(), combinations.combine(reshuffle=[True, False]), combinations.combine(graph_seed=38, op_seed=None) + combinations.combine(graph_seed=None, op_seed=42) + @@ -188,7 +188,7 @@ class ShuffleTest(test_base.DatasetTestBase, parameterized.TestCase): # TODO(b/117581999): enable this test for eager-mode. @combinations.generate( combinations.times( - combinations.combine(tf_api_version=[1, 2], mode="graph"), + test_base.graph_only_combinations(), combinations.combine( reshuffle=[True, False], initializable=[True, False]))) def testMultipleIterators(self, reshuffle, initializable): @@ -278,7 +278,7 @@ class ShuffleTest(test_base.DatasetTestBase, parameterized.TestCase): @combinations.generate( combinations.times( - combinations.combine(tf_api_version=[1, 2], mode="eager"), + test_base.eager_only_combinations(), combinations.combine(reshuffle=[True, False], seed=[None, 42]))) def testReshuffleSeparateTransformations(self, reshuffle, seed): dataset = dataset_ops.Dataset.range(10) diff --git a/tensorflow/python/data/kernel_tests/skip_test.py b/tensorflow/python/data/kernel_tests/skip_test.py index 74dc8b7f55c..176893d90d2 100644 --- a/tensorflow/python/data/kernel_tests/skip_test.py +++ b/tensorflow/python/data/kernel_tests/skip_test.py @@ -17,46 +17,30 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.framework import test_util +from tensorflow.python.framework import combinations from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes -class SkipTest(test_base.DatasetTestBase): +class SkipTest(test_base.DatasetTestBase, parameterized.TestCase): - def testSkipTensorDataset(self): + @combinations.generate( + combinations.times(test_base.default_test_combinations(), + combinations.combine(count=[-1, 0, 4, 10, 25]))) + def testBasic(self, count): components = (np.arange(10),) - - def do_test(count): - dataset = dataset_ops.Dataset.from_tensor_slices(components).skip(count) - self.assertEqual( - [c.shape[1:] for c in components], - [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) - start_range = min(count, 10) if count != -1 else 10 - self.assertDatasetProduces( - dataset, - [tuple(components[0][i:i + 1]) for i in range(start_range, 10)]) - - # Skip fewer than input size, we should skip - # the first 4 elements and then read the rest. - do_test(4) - - # Skip more than input size: get nothing. - do_test(25) - - # Skip exactly input size. - do_test(10) - - # Set -1 for 'count': skip the entire dataset. - do_test(-1) - - # Skip nothing - do_test(0) - + dataset = dataset_ops.Dataset.from_tensor_slices(components).skip(count) + self.assertEqual( + [c.shape[1:] for c in components], + [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) + start_range = min(count, 10) if count != -1 else 10 + self.assertDatasetProduces( + dataset, + [tuple(components[0][i:i + 1]) for i in range(start_range, 10)]) if __name__ == "__main__": diff --git a/tensorflow/python/data/kernel_tests/take_test.py b/tensorflow/python/data/kernel_tests/take_test.py index 665ed59a7bc..14796551e16 100644 --- a/tensorflow/python/data/kernel_tests/take_test.py +++ b/tensorflow/python/data/kernel_tests/take_test.py @@ -17,40 +17,30 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.framework import test_util +from tensorflow.python.framework import combinations from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes -class TakeTest(test_base.DatasetTestBase): +class TakeTest(test_base.DatasetTestBase, parameterized.TestCase): - def testTakeTensorDataset(self): + @combinations.generate( + combinations.times(test_base.default_test_combinations(), + combinations.combine(count=[-1, 0, 4, 10, 25]))) + def testBasic(self, count): components = (np.arange(10),) + dataset = dataset_ops.Dataset.from_tensor_slices(components).take(count) + self.assertEqual( + [c.shape[1:] for c in components], + [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) + num_output = min(count, 10) if count != -1 else 10 + self.assertDatasetProduces( + dataset, [tuple(components[0][i:i + 1]) for i in range(num_output)]) - def do_test(count): - dataset = dataset_ops.Dataset.from_tensor_slices(components).take(count) - self.assertEqual( - [c.shape[1:] for c in components], - [shape for shape in dataset_ops.get_legacy_output_shapes(dataset)]) - num_output = min(count, 10) if count != -1 else 10 - self.assertDatasetProduces( - dataset, [tuple(components[0][i:i + 1]) for i in range(num_output)]) - - # Take fewer than input size - do_test(4) - - # Take more than input size - do_test(25) - - # Take all of input - do_test(-1) - - # Take nothing - do_test(0) if __name__ == "__main__": test.main() diff --git a/tensorflow/python/data/kernel_tests/test_base.py b/tensorflow/python/data/kernel_tests/test_base.py index 6dfee4cc0f7..60796b178bf 100644 --- a/tensorflow/python/data/kernel_tests/test_base.py +++ b/tensorflow/python/data/kernel_tests/test_base.py @@ -58,7 +58,11 @@ class DatasetTestBase(test.TestCase): def assertValuesEqual(self, expected, actual): """Asserts that two values are equal.""" - if sparse_tensor.is_sparse(expected): + if isinstance(expected, dict): + self.assertItemsEqual(list(expected.keys()), list(actual.keys())) + for k in expected.keys(): + self.assertValuesEqual(expected[k], actual[k]) + elif sparse_tensor.is_sparse(expected): self.assertAllEqual(expected.indices, actual.indices) self.assertAllEqual(expected.values, actual.values) self.assertAllEqual(expected.dense_shape, actual.dense_shape) diff --git a/tensorflow/python/data/kernel_tests/text_line_dataset_test.py b/tensorflow/python/data/kernel_tests/text_line_dataset_test.py index c62d4ec8270..35b479faa21 100644 --- a/tensorflow/python/data/kernel_tests/text_line_dataset_test.py +++ b/tensorflow/python/data/kernel_tests/text_line_dataset_test.py @@ -21,11 +21,12 @@ import gzip import os import zlib +from absl.testing import parameterized + from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import readers -from tensorflow.python.eager import context -from tensorflow.python.framework import test_util +from tensorflow.python.framework import combinations from tensorflow.python.platform import test from tensorflow.python.util import compat @@ -37,8 +38,7 @@ except ImportError: psutil_import_succeeded = False -@test_util.run_all_in_graph_and_eager_modes -class TextLineDatasetTest(test_base.DatasetTestBase): +class TextLineDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): def _lineText(self, f, l): return compat.as_bytes("%d: %d" % (f, l)) @@ -76,7 +76,11 @@ class TextLineDatasetTest(test_base.DatasetTestBase): return filenames - def _testTextLineDataset(self, compression_type=None): + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine(compression_type=[None, "GZIP", "ZLIB"]))) + def testTextLineDataset(self, compression_type): test_filenames = self._createFiles( 2, 5, crlf=True, compression_type=compression_type) @@ -115,6 +119,7 @@ class TextLineDatasetTest(test_base.DatasetTestBase): expected_output=[[self._lineText(0, i) for i in range(5)], [self._lineText(1, i) for i in range(5)]] * 10) + @combinations.generate(test_base.default_test_combinations()) def testTextLineDatasetParallelRead(self): test_filenames = self._createFiles(10, 10) files = dataset_ops.Dataset.from_tensor_slices(test_filenames).repeat(10) @@ -125,15 +130,7 @@ class TextLineDatasetTest(test_base.DatasetTestBase): self.assertDatasetProduces( dataset, expected_output=expected_output * 10, assert_items_equal=True) - def testTextLineDatasetNoCompression(self): - self._testTextLineDataset() - - def testTextLineDatasetGzipCompression(self): - self._testTextLineDataset(compression_type="GZIP") - - def testTextLineDatasetZlibCompression(self): - self._testTextLineDataset(compression_type="ZLIB") - + @combinations.generate(test_base.default_test_combinations()) def testTextLineDatasetBuffering(self): test_filenames = self._createFiles(2, 5, crlf=True) @@ -143,33 +140,33 @@ class TextLineDatasetTest(test_base.DatasetTestBase): expected_output.extend([self._lineText(j, i) for i in range(5)]) self.assertDatasetProduces(repeat_dataset, expected_output=expected_output) + @combinations.generate(test_base.eager_only_combinations()) def testIteratorResourceCleanup(self): filename = os.path.join(self.get_temp_dir(), "text.txt") with open(filename, "wt") as f: for i in range(3): f.write("%d\n" % (i,)) - with context.eager_mode(): - first_iterator = iter(readers.TextLineDataset(filename)) - self.assertEqual(b"0", next(first_iterator).numpy()) - second_iterator = iter(readers.TextLineDataset(filename)) - self.assertEqual(b"0", next(second_iterator).numpy()) - # Eager kernel caching is based on op attributes, which includes the - # Dataset's output shape. Create a different kernel to test that they - # don't create resources with the same names. - different_kernel_iterator = iter( - readers.TextLineDataset(filename).repeat().batch(16)) - self.assertEqual([16], next(different_kernel_iterator).shape) - # Remove our references to the Python Iterator objects, which (assuming no - # reference cycles) is enough to trigger DestroyResourceOp and close the - # partially-read files. - del first_iterator - del second_iterator - del different_kernel_iterator - if not psutil_import_succeeded: - self.skipTest( - "psutil is required to check that we've closed our files.") - open_files = psutil.Process().open_files() - self.assertNotIn(filename, [open_file.path for open_file in open_files]) + first_iterator = iter(readers.TextLineDataset(filename)) + self.assertEqual(b"0", next(first_iterator).numpy()) + second_iterator = iter(readers.TextLineDataset(filename)) + self.assertEqual(b"0", next(second_iterator).numpy()) + # Eager kernel caching is based on op attributes, which includes the + # Dataset's output shape. Create a different kernel to test that they + # don't create resources with the same names. + different_kernel_iterator = iter( + readers.TextLineDataset(filename).repeat().batch(16)) + self.assertEqual([16], next(different_kernel_iterator).shape) + # Remove our references to the Python Iterator objects, which (assuming no + # reference cycles) is enough to trigger DestroyResourceOp and close the + # partially-read files. + del first_iterator + del second_iterator + del different_kernel_iterator + if not psutil_import_succeeded: + self.skipTest( + "psutil is required to check that we've closed our files.") + open_files = psutil.Process().open_files() + self.assertNotIn(filename, [open_file.path for open_file in open_files]) if __name__ == "__main__": diff --git a/tensorflow/python/data/kernel_tests/tf_record_dataset_test.py b/tensorflow/python/data/kernel_tests/tf_record_dataset_test.py index 5cf8308a55f..792c4926640 100644 --- a/tensorflow/python/data/kernel_tests/tf_record_dataset_test.py +++ b/tensorflow/python/data/kernel_tests/tf_record_dataset_test.py @@ -21,31 +21,31 @@ import gzip import os import zlib +from absl.testing import parameterized + from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import readers +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op -from tensorflow.python.framework import test_util from tensorflow.python.lib.io import python_io from tensorflow.python.platform import test from tensorflow.python.util import compat -@test_util.run_all_in_graph_and_eager_modes -class TFRecordDatasetTest(test_base.DatasetTestBase): +class TFRecordDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): def setUp(self): super(TFRecordDatasetTest, self).setUp() self._num_files = 2 self._num_records = 7 - self.test_filenames = self._createFiles() - def dataset_fn(self, - filenames, - compression_type="", - num_epochs=1, - batch_size=None): + def _dataset_factory(self, + filenames, + compression_type="", + num_epochs=1, + batch_size=None): repeat_dataset = readers.TFRecordDataset( filenames, compression_type).repeat(num_epochs) @@ -67,6 +67,7 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): writer.close() return filenames + @combinations.generate(test_base.default_test_combinations()) def testTFRecordDatasetConstructorErrorsTensorInput(self): with self.assertRaisesRegex(TypeError, "filenames.*must be.*Tensor.*string"): @@ -78,37 +79,40 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): with self.assertRaises(Exception): readers.TFRecordDataset(object()) + @combinations.generate(test_base.default_test_combinations()) def testReadOneEpoch(self): # Basic test: read from file 0. - dataset = self.dataset_fn(self.test_filenames[0]) + dataset = self._dataset_factory(self.test_filenames[0]) self.assertDatasetProduces( dataset, expected_output=[self._record(0, i) for i in range(self._num_records)]) # Basic test: read from file 1. - dataset = self.dataset_fn(self.test_filenames[1]) + dataset = self._dataset_factory(self.test_filenames[1]) self.assertDatasetProduces( dataset, expected_output=[self._record(1, i) for i in range(self._num_records)]) # Basic test: read from both files. - dataset = self.dataset_fn(self.test_filenames) + dataset = self._dataset_factory(self.test_filenames) expected_output = [] for j in range(self._num_files): expected_output.extend( [self._record(j, i) for i in range(self._num_records)]) self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testReadTenEpochs(self): - dataset = self.dataset_fn(self.test_filenames, num_epochs=10) + dataset = self._dataset_factory(self.test_filenames, num_epochs=10) expected_output = [] for j in range(self._num_files): expected_output.extend( [self._record(j, i) for i in range(self._num_records)]) self.assertDatasetProduces(dataset, expected_output=expected_output * 10) + @combinations.generate(test_base.default_test_combinations()) def testReadTenEpochsOfBatches(self): - dataset = self.dataset_fn( + dataset = self._dataset_factory( self.test_filenames, num_epochs=10, batch_size=self._num_records) expected_output = [] for j in range(self._num_files): @@ -116,6 +120,7 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): [self._record(j, i) for i in range(self._num_records)]) self.assertDatasetProduces(dataset, expected_output=expected_output * 10) + @combinations.generate(test_base.default_test_combinations()) def testReadZlibFiles(self): zlib_files = [] for i, fn in enumerate(self.test_filenames): @@ -130,9 +135,10 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): for j in range(self._num_files): expected_output.extend( [self._record(j, i) for i in range(self._num_records)]) - dataset = self.dataset_fn(zlib_files, compression_type="ZLIB") + dataset = self._dataset_factory(zlib_files, compression_type="ZLIB") self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testReadGzipFiles(self): gzip_files = [] for i, fn in enumerate(self.test_filenames): @@ -145,9 +151,10 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): for j in range(self._num_files): expected_output.extend( [self._record(j, i) for i in range(self._num_records)]) - dataset = self.dataset_fn(gzip_files, compression_type="GZIP") + dataset = self._dataset_factory(gzip_files, compression_type="GZIP") self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testReadWithBuffer(self): one_mebibyte = 2**20 dataset = readers.TFRecordDataset( @@ -158,6 +165,7 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): [self._record(j, i) for i in range(self._num_records)]) self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testReadFromDatasetOfFiles(self): files = dataset_ops.Dataset.from_tensor_slices(self.test_filenames) expected_output = [] @@ -167,6 +175,7 @@ class TFRecordDatasetTest(test_base.DatasetTestBase): dataset = readers.TFRecordDataset(files) self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testReadTenEpochsFromDatasetOfFilesInParallel(self): files = dataset_ops.Dataset.from_tensor_slices( self.test_filenames).repeat(10) diff --git a/tensorflow/python/data/kernel_tests/unbatch_test.py b/tensorflow/python/data/kernel_tests/unbatch_test.py index 5bb4852d534..44d949385b0 100644 --- a/tensorflow/python/data/kernel_tests/unbatch_test.py +++ b/tensorflow/python/data/kernel_tests/unbatch_test.py @@ -23,11 +23,11 @@ import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor -from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import string_ops @@ -36,13 +36,14 @@ from tensorflow.python.platform import test from tensorflow.python.util import compat -@test_util.run_all_in_graph_and_eager_modes class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): + @combinations.generate(test_base.default_test_combinations()) def testUnbatchWithUnknownRankInput(self): dataset = dataset_ops.Dataset.from_tensors([0, 1, 2, 3]).unbatch() self.assertDatasetProduces(dataset, range(4)) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchScalarDataset(self): data = tuple([math_ops.range(10) for _ in range(3)]) data = dataset_ops.Dataset.from_tensor_slices(data) @@ -54,12 +55,14 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces(data, [(i,) * 3 for i in range(10)]) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchNestedDataset(self): data = dataset_ops.Dataset.from_tensors( [dataset_ops.Dataset.range(10) for _ in range(10)]) data = data.unbatch().flat_map(lambda x: x) self.assertDatasetProduces(data, list(range(10)) * 10) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchDatasetWithStrings(self): data = tuple([math_ops.range(10) for _ in range(3)]) data = dataset_ops.Dataset.from_tensor_slices(data) @@ -73,6 +76,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( data, [(i, compat.as_bytes(str(i)), i) for i in range(10)]) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchDatasetWithSparseTensor(self): st = sparse_tensor.SparseTensorValue( indices=[[i, i] for i in range(10)], @@ -87,6 +91,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): ] self.assertDatasetProduces(data, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchDatasetWithDenseSparseAndRaggedTensor(self): st = sparse_tensor.SparseTensorValue( indices=[[i, i] for i in range(10)], @@ -104,6 +109,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( data, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchDatasetWithRaggedTensor(self): rt = ragged_factory_ops.constant_value([[[0]], [[1]], [[2]], [[3]], [[4]], [[5]], [[6]], [[7]], [[8]], [[9]]]) @@ -119,6 +125,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces( data, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchSingleElementTupleDataset(self): data = tuple([(math_ops.range(10),) for _ in range(3)]) data = dataset_ops.Dataset.from_tensor_slices(data) @@ -130,6 +137,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): self.assertDatasetProduces(data, [((i,),) * 3 for i in range(10)]) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchMultiElementTupleDataset(self): data = tuple([(math_ops.range(10 * i, 10 * i + 10), array_ops.fill([10], "hi")) for i in range(3)]) @@ -146,6 +154,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): data, [((i, b"hi"), (10 + i, b"hi"), (20 + i, b"hi")) for i in range(10)]) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchEmpty(self): data = dataset_ops.Dataset.from_tensors( (constant_op.constant([]), constant_op.constant([], shape=[0, 4]), @@ -153,15 +162,15 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): data = data.unbatch() self.assertDatasetProduces(data, []) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchStaticShapeMismatch(self): data = dataset_ops.Dataset.from_tensors((np.arange(7), np.arange(8), np.arange(9))) with self.assertRaises(ValueError): data.unbatch() - # Note: dynamic shape mismatch is graph specific test. - @test_util.run_deprecated_v1 - def testSkipEagerUnbatchDynamicShapeMismatch(self): + @combinations.generate(test_base.graph_only_combinations()) + def testUnbatchDynamicShapeMismatch(self): ph1 = array_ops.placeholder(dtypes.int32, shape=[None]) ph2 = array_ops.placeholder(dtypes.int32, shape=None) data = dataset_ops.Dataset.from_tensors((ph1, ph2)) @@ -190,6 +199,7 @@ class UnbatchTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(next_element) + @combinations.generate(test_base.default_test_combinations()) def testUnbatchDatasetWithUintDtypes(self): components = ( np.tile(np.array([[0], [1], [2], [3]], dtype=np.uint8), 2), diff --git a/tensorflow/python/data/kernel_tests/window_test.py b/tensorflow/python/data/kernel_tests/window_test.py index 122e874f0a0..98b453a5900 100644 --- a/tensorflow/python/data/kernel_tests/window_test.py +++ b/tensorflow/python/data/kernel_tests/window_test.py @@ -24,43 +24,32 @@ from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest from tensorflow.python.eager import context +from tensorflow.python.framework import combinations from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import sparse_tensor -from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): - @parameterized.named_parameters( - ("1", 20, 14, 7, 1), - ("2", 20, 17, 9, 1), - ("3", 20, 14, 14, 1), - ("4", 20, 10, 14, 1), - ("5", 20, 14, 19, 1), - ("6", 20, 4, 1, 2), - ("7", 20, 2, 1, 6), - ("8", 20, 4, 7, 2), - ("9", 20, 2, 7, 6), - ("10", 1, 10, 4, 1), - ("11", 0, 10, 4, 1), - ("12", 20, 14, 7, 1, False), - ("13", 20, 17, 9, 1, False), - ("14", 20, 14, 14, 1, False), - ("15", 20, 10, 14, 1, False), - ("16", 20, 14, 19, 1, False), - ("17", 20, 4, 1, 2, False), - ("18", 20, 2, 1, 6, False), - ("19", 20, 4, 7, 2, False), - ("20", 20, 2, 7, 6, False), - ("21", 1, 10, 4, 1, False), - ("22", 0, 10, 4, 1, False), - ) - def testWindowDataset(self, count, size, shift, stride, drop_remainder=True): + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine( + count=20, + size=[10, 14, 17], + shift=[7, 14], + stride=[1, 2, 6], + drop_remainder=[True, False]) + combinations.combine( + count=[0, 1], + size=10, + shift=4, + stride=1, + drop_remainder=[True, False]))) + def testWindowDataset(self, count, size, shift, stride, drop_remainder): """Tests a dataset that slides a window its input elements.""" components = (np.arange(7), np.array([[1, 2, 3]]) * np.arange(7)[:, np.newaxis], @@ -111,11 +100,12 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - @parameterized.named_parameters( - ("1", 14, 0, 3, 1), - ("2", 14, 3, 0, 1), - ("3", 14, 3, 3, 0), - ) + @combinations.generate( + combinations.times( + test_base.default_test_combinations(), + combinations.combine(count=20, size=0, shift=3, stride=1) + + combinations.combine(count=20, size=3, shift=0, stride=1) + + combinations.combine(count=20, size=3, shift=3, stride=0))) def testWindowDatasetInvalid(self, count, size, shift, stride): with self.assertRaises(errors.InvalidArgumentError): ds = dataset_ops.Dataset.range(10).map(lambda x: x).repeat(count).window( @@ -123,12 +113,14 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): stride=stride).flat_map(lambda x: x.batch(batch_size=size)) self.evaluate(ds._variant_tensor) + @combinations.generate(test_base.default_test_combinations()) def testWindowDifferentNestedStructures(self): ds = dataset_ops.Dataset.from_tensor_slices(([1, 2], [3, 4])).window(2) self.getNext(ds) ds = dataset_ops.Dataset.from_tensor_slices({"a": [1, 2]}).window(2) self.getNext(ds) + @combinations.generate(test_base.default_test_combinations()) def testWindowSparse(self): def _sparse(i): @@ -148,6 +140,7 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): ] self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testWindowSparseWithDifferentDenseShapes(self): def _sparse(i): @@ -177,6 +170,7 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): dense_shape=[5, i * 3 + 5 - 1])) self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testNestedWindowSparse(self): def _sparse(i): @@ -205,6 +199,7 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): ] self.assertDatasetProduces(dataset, expected_output=expected_output) + @combinations.generate(test_base.default_test_combinations()) def testWindowShapeError(self): def generator(): @@ -222,6 +217,7 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): r"Cannot batch tensors with different shapes in component 0. " r"First element had shape \[3\] and element 2 had shape \[4\].")) + @combinations.generate(test_base.default_test_combinations()) def testWindowIgnoreErrors(self): input_values = np.float32([1., np.nan, 2., np.nan, 3.]) dataset = dataset_ops.Dataset.from_tensor_slices(input_values).map( @@ -232,6 +228,7 @@ class WindowTest(test_base.DatasetTestBase, parameterized.TestCase): dataset, expected_output=[np.float32([1., 2.]), np.float32([2., 3.])]) + @combinations.generate(test_base.default_test_combinations()) def testNestedOutput(self): if not context.executing_eagerly(): self.skipTest("self.evaluate() does not work with a dataset") diff --git a/tensorflow/python/data/kernel_tests/zip_test.py b/tensorflow/python/data/kernel_tests/zip_test.py index 72f739e4e4e..c63091754c3 100644 --- a/tensorflow/python/data/kernel_tests/zip_test.py +++ b/tensorflow/python/data/kernel_tests/zip_test.py @@ -17,66 +17,68 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from absl.testing import parameterized import numpy as np from tensorflow.python.data.kernel_tests import test_base from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.framework import combinations from tensorflow.python.framework import errors from tensorflow.python.framework import tensor_shape -from tensorflow.python.framework import test_util from tensorflow.python.platform import test -@test_util.run_all_in_graph_and_eager_modes -class ZipTest(test_base.DatasetTestBase): +def _dataset_factory(components): + datasets = tuple([ + dataset_ops.Dataset.from_tensor_slices(component) + for component in components + ]) + return dataset_ops.Dataset.zip(datasets) - def testZipDataset(self): - def dataset_fn(components): - datasets = tuple([ - dataset_ops.Dataset.from_tensor_slices(component) - for component in components - ]) - return dataset_ops.Dataset.zip(datasets) +class ZipTest(test_base.DatasetTestBase, parameterized.TestCase): - equal_length_components = [ + @combinations.generate(test_base.default_test_combinations()) + def testZipEqual(self): + components = [ np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile(np.array([[12], [13], [14], [15]]), 22), np.array([37.0, 38.0, 39.0, 40.0]) ] - - get_next = self.getNext(dataset_fn(equal_length_components)) + get_next = self.getNext(_dataset_factory(components)) for i in range(4): results = self.evaluate(get_next()) - for component, result_component in zip(equal_length_components, results): + for component, result_component in zip(components, results): self.assertAllEqual(component[i], result_component) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - variable_length_components = [[1, 2, 3, 4], [1, 2, 3, 4, 5], [1.0, 2.0]] - get_next = self.getNext(dataset_fn(variable_length_components)) + @combinations.generate(test_base.default_test_combinations()) + def testZipUnequal(self): + components = [[1, 2, 3, 4], [1, 2, 3, 4, 5], [1.0, 2.0]] + get_next = self.getNext(_dataset_factory(components)) for i in range(2): results = self.evaluate(get_next()) - for component, result_component in zip(variable_length_components, - results): + for component, result_component in zip(components, results): self.assertAllEqual(component[i], result_component) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) - def testNestedZipDataset(self): + @combinations.generate(test_base.default_test_combinations()) + def testNested(self): - equal_length_components = [ + components = [ np.tile(np.array([[1], [2], [3], [4]]), 20), np.tile(np.array([[12], [13], [14], [15]]), 22), np.array([37.0, 38.0, 39.0, 40.0]) ] datasets = [ dataset_ops.Dataset.from_tensor_slices(component) - for component in equal_length_components + for component in components ] dataset = dataset_ops.Dataset.zip((datasets[0], (datasets[1], datasets[2]))) @@ -88,9 +90,9 @@ class ZipTest(test_base.DatasetTestBase): get_next = self.getNext(dataset) for i in range(4): result1, (result2, result3) = self.evaluate(get_next()) - self.assertAllEqual(equal_length_components[0][i], result1) - self.assertAllEqual(equal_length_components[1][i], result2) - self.assertAllEqual(equal_length_components[2][i], result3) + self.assertAllEqual(components[0][i], result1) + self.assertAllEqual(components[1][i], result2) + self.assertAllEqual(components[2][i], result3) with self.assertRaises(errors.OutOfRangeError): self.evaluate(get_next()) with self.assertRaises(errors.OutOfRangeError): diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index f67dec9d720..f3367023a7b 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -66,7 +66,6 @@ from tensorflow.python.ops import gen_io_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import script_ops from tensorflow.python.ops import string_ops -from tensorflow.python.platform import tf_logging as logging from tensorflow.python.training.tracking import base as tracking_base from tensorflow.python.training.tracking import tracking from tensorflow.python.util import deprecation @@ -2437,26 +2436,25 @@ class DatasetV1Adapter(DatasetV1): def _ensure_same_dataset_graph(dataset): """Walks the dataset graph to ensure all datasets come from the same graph.""" + # pylint: disable=protected-access current_graph = ops.get_default_graph() bfs_q = Queue.Queue() - bfs_q.put(dataset) # pylint: disable=protected-access + bfs_q.put(dataset) visited = [] while not bfs_q.empty(): ds = bfs_q.get() visited.append(ds) - ds_graph = ds._graph # pylint: disable=protected-access + ds_graph = ds._graph if current_graph != ds_graph: - logging.warning("The graph (" + str(current_graph) + ") of the iterator " - "is different from the graph (" + str(ds_graph) + ") " - "the dataset: " + str(ds._variant_tensor) + " was " # pylint: disable=protected-access - "created in. If you are using the Estimator API, " - "make sure that no part of the dataset returned by the " - "`input_fn` function is defined outside the `input_fn` " - "function. Please ensure that all datasets in the " - "pipeline are created in the same graph as the iterator. " - "NOTE: This warning will become an error in future " - "versions of TensorFlow.") - for input_ds in ds._inputs(): # pylint: disable=protected-access + raise ValueError( + "The graph (" + str(current_graph) + ") of the iterator is different " + "from the graph (" + str(ds_graph) + ") the dataset: " + + str(ds._variant_tensor) + " was created in. If you are using the " + "Estimator API, make sure that no part of the dataset returned by " + "the `input_fn` function is defined outside the `input_fn` function. " + "Please ensure that all datasets in the pipeline are created in the " + "same graph as the iterator.") + for input_ds in ds._inputs(): if input_ds not in visited: bfs_q.put(input_ds)