diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/unique_op_test.py')
-rw-r--r-- | tensorflow/python/kernel_tests/unique_op_test.py | 26 |
1 files changed, 25 insertions, 1 deletions
diff --git a/tensorflow/python/kernel_tests/unique_op_test.py b/tensorflow/python/kernel_tests/unique_op_test.py index a1903887c7..a50f53b3cd 100644 --- a/tensorflow/python/kernel_tests/unique_op_test.py +++ b/tensorflow/python/kernel_tests/unique_op_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.platform import test @@ -37,6 +38,17 @@ class UniqueTest(test.TestCase): for i in range(len(x)): self.assertEqual(x[i], tf_y[tf_idx[i]]) + def testInt32OutIdxInt64(self): + x = np.random.randint(2, high=10, size=7000) + with self.test_session() as sess: + y, idx = array_ops.unique(x, out_idx=dtypes.int64) + tf_y, tf_idx = sess.run([y, idx]) + + self.assertEqual(len(x), len(tf_idx)) + self.assertEqual(len(tf_y), len(np.unique(x))) + for i in range(len(x)): + self.assertEqual(x[i], tf_y[tf_idx[i]]) + def testString(self): indx = np.random.randint(65, high=122, size=7000) x = [chr(i) for i in indx] @@ -49,7 +61,6 @@ class UniqueTest(test.TestCase): for i in range(len(x)): self.assertEqual(x[i], tf_y[tf_idx[i]].decode('ascii')) - class UniqueWithCountsTest(test.TestCase): def testInt32(self): @@ -65,6 +76,19 @@ class UniqueWithCountsTest(test.TestCase): for value, count in zip(tf_y, tf_count): self.assertEqual(count, np.sum(x == value)) + def testInt32OutIdxInt64(self): + x = np.random.randint(2, high=10, size=7000) + with self.test_session() as sess: + y, idx, count = array_ops.unique_with_counts(x, out_idx=dtypes.int64) + tf_y, tf_idx, tf_count = sess.run([y, idx, count]) + + self.assertEqual(len(x), len(tf_idx)) + self.assertEqual(len(tf_y), len(np.unique(x))) + for i in range(len(x)): + self.assertEqual(x[i], tf_y[tf_idx[i]]) + for value, count in zip(tf_y, tf_count): + self.assertEqual(count, np.sum(x == value)) + def testString(self): indx = np.random.randint(65, high=122, size=7000) x = [chr(i) for i in indx] |