Add test for bug in CUB that caused dynamic partition to fail on the GPU.

PiperOrigin-RevId: 186834668

1 files changed, 12 insertions, 0 deletions

diff --git a/tensorflow/python/kernel_tests/dynamic_partition_op_test.py b/tensorflow/python/kernel_tests/dynamic_partition_op_test.py
index fedbf9e696..5e8937ad2c 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()