diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/bincount_op_test.py')
-rw-r--r-- | tensorflow/python/kernel_tests/bincount_op_test.py | 25 |
1 files changed, 18 insertions, 7 deletions
diff --git a/tensorflow/python/kernel_tests/bincount_op_test.py b/tensorflow/python/kernel_tests/bincount_op_test.py index 7a610debd1..79285476b4 100644 --- a/tensorflow/python/kernel_tests/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/bincount_op_test.py @@ -25,11 +25,10 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops import math_ops from tensorflow.python.platform import googletest - class BincountTest(test_util.TensorFlowTestCase): def test_empty(self): - with self.test_session(): + with self.test_session(use_gpu=True): self.assertAllEqual( math_ops.bincount([], minlength=5).eval(), [0, 0, 0, 0, 0]) self.assertAllEqual(math_ops.bincount([], minlength=1).eval(), [0]) @@ -42,7 +41,7 @@ class BincountTest(test_util.TensorFlowTestCase): np.float64) def test_values(self): - with self.test_session(): + with self.test_session(use_gpu=True): self.assertAllEqual( math_ops.bincount([1, 1, 1, 2, 2, 3]).eval(), [0, 3, 2, 1]) arr = [1, 1, 2, 1, 2, 3, 1, 2, 3, 4, 1, 2, 3, 4, 5] @@ -57,14 +56,14 @@ class BincountTest(test_util.TensorFlowTestCase): math_ops.bincount(np.arange(10000)).eval(), np.ones(10000)) def test_maxlength(self): - with self.test_session(): + with self.test_session(use_gpu=True): self.assertAllEqual(math_ops.bincount([5], maxlength=3).eval(), [0, 0, 0]) self.assertAllEqual(math_ops.bincount([1], maxlength=3).eval(), [0, 1]) self.assertAllEqual(math_ops.bincount([], maxlength=3).eval(), []) def test_random_with_weights(self): num_samples = 10000 - with self.test_session(): + with self.test_session(use_gpu=True): np.random.seed(42) for dtype in [dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64]: arr = np.random.randint(0, 1000, num_samples) @@ -72,17 +71,29 @@ class BincountTest(test_util.TensorFlowTestCase): weights = np.random.randint(-100, 100, num_samples) else: weights = np.random.random(num_samples) - self.assertAllEqual( + self.assertAllClose( math_ops.bincount(arr, weights).eval(), np.bincount(arr, weights)) + def test_random_without_weights(self): + num_samples = 10000 + with self.test_session(use_gpu=True): + np.random.seed(42) + for dtype in [np.int32, np.float32]: + arr = np.random.randint(0, 1000, num_samples) + weights = np.ones(num_samples).astype(dtype) + self.assertAllClose( + math_ops.bincount(arr, None).eval(), + np.bincount(arr, weights)) + def test_zero_weights(self): - with self.test_session(): + with self.test_session(use_gpu=True): self.assertAllEqual( math_ops.bincount(np.arange(1000), np.zeros(1000)).eval(), np.zeros(1000)) def test_negative(self): + # unsorted_segment_sum will only report InvalidArgumentError on CPU with self.test_session(): with self.assertRaises(errors.InvalidArgumentError): math_ops.bincount([1, 2, 3, -1, 6, 8]).eval() |