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, 26 insertions, 0 deletions
diff --git a/tensorflow/python/kernel_tests/unique_op_test.py b/tensorflow/python/kernel_tests/unique_op_test.py index a50f53b3cd..04758ce45a 100644 --- a/tensorflow/python/kernel_tests/unique_op_test.py +++ b/tensorflow/python/kernel_tests/unique_op_test.py @@ -22,6 +22,7 @@ import numpy as np from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_array_ops from tensorflow.python.platform import test @@ -61,6 +62,31 @@ class UniqueTest(test.TestCase): for i in range(len(x)): self.assertEqual(x[i], tf_y[tf_idx[i]].decode('ascii')) + def testInt32Axis(self): + x = np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]]) + with self.test_session() as sess: + y0, idx0 = gen_array_ops.unique_v2(x, axis=[0]) + tf_y0, tf_idx0 = sess.run([y0, idx0]) + y1, idx1 = gen_array_ops.unique_v2(x, axis=[1]) + tf_y1, tf_idx1 = sess.run([y1, idx1]) + self.assertAllEqual(tf_y0, np.array([[1, 0, 0], [2, 0, 0]])) + self.assertAllEqual(tf_idx0, np.array([0, 0, 1])) + self.assertAllEqual(tf_y1, np.array([[1, 0], [1, 0], [2, 0]])) + self.assertAllEqual(tf_idx1, np.array([0, 1, 1])) + + def testInt32V2(self): + # This test is only temporary, once V2 is used + # by default, the axis will be wrapped to allow `axis=None`. + x = np.random.randint(2, high=10, size=7000) + with self.test_session() as sess: + y, idx = gen_array_ops.unique_v2(x, axis=[]) + 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]]) + class UniqueWithCountsTest(test.TestCase): def testInt32(self): |