diff options
Diffstat (limited to 'tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py')
-rw-r--r-- | tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py | 25 |
1 files changed, 1 insertions, 24 deletions
diff --git a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py index 6001721726..347af18576 100644 --- a/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py +++ b/tensorflow/python/data/kernel_tests/shuffle_dataset_op_test.py @@ -19,7 +19,6 @@ from __future__ import print_function import collections -from absl.testing import parameterized import numpy as np from tensorflow.python.data.kernel_tests import test_base @@ -32,7 +31,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.platform import test -class ShuffleDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): +class ShuffleDatasetTest(test_base.DatasetTestBase): def testShuffleDataset(self): components = ( @@ -210,27 +209,5 @@ class ShuffleDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): with self.assertRaises(errors.OutOfRangeError): sess.run(next_element) - @parameterized.named_parameters( - ("ReshuffleEachIterationNoSeed", None, True), - ("ReshuffleEachIterationWithSeed", 42, True), - ("NoReshuffleEachIterationNoSeed", None, False), - ("NoReshuffleEachIterationWithSeed", 42, False), - ) - def testShuffleAndZipDataset(self, seed, reshuffle): - dataset = (dataset_ops.Dataset.range(10) - .shuffle(10, seed=seed, reshuffle_each_iteration=reshuffle) - .repeat(3)) - dataset = dataset_ops.Dataset.zip((dataset, dataset)) - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - - with self.cached_session() as sess: - for _ in range(30): - x, y = sess.run(next_element) - self.assertEqual(x, y) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - if __name__ == "__main__": test.main() |