From cc171bb7371590ee45e361b6a50a018d026412f6 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 23 Feb 2018 15:43:09 -0800 Subject: [PATCH] Add test for bug in CUB that caused dynamic partition to fail on the GPU. PiperOrigin-RevId: 186834668 --- tensorflow/python/kernel_tests/dynamic_partition_op_test.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorflow/python/kernel_tests/dynamic_partition_op_test.py b/tensorflow/python/kernel_tests/dynamic_partition_op_test.py index fedbf9e..5e8937a 100644 --- a/tensorflow/python/kernel_tests/dynamic_partition_op_test.py +++ b/tensorflow/python/kernel_tests/dynamic_partition_op_test.py @@ -326,6 +326,18 @@ class DynamicPartitionTest(test.TestCase): with self.assertRaises(ValueError): data_flow_ops.dynamic_partition(data, indices, num_partitions=4) + # see https://github.com/tensorflow/tensorflow/issues/17106 + def testCUBBug(self): + x = constant_op.constant(np.random.randn(3072)) + inds = [0]*189 + [1]*184 + [2]*184 + [3]*191 + [4]*192 + [5]*195 + [6]*195 + inds += [7]*195 + [8]*188 + [9]*195 + [10]*188 + [11]*202 + [12]*194 + inds += [13]*194 + [14]*194 + [15]*192 + self.assertEqual(len(inds), x.shape[0]) + partitioned = data_flow_ops.dynamic_partition(x, inds, 16) + with self.test_session() as sess: + res = sess.run(partitioned) + self.assertEqual(res[-1].shape[0], 192) + if __name__ == "__main__": test.main() -- 2.7.4