Fix crash in random.shuffle operator (apache#15041)

* fix crash in random_shuffle caused by int overflow * add unit test * add comment * remove small random test to avoid CI failure
haohuanw · Jun 23, 2019 · 5abe119 · 5abe119
1 parent 5798abf
commit 5abe119
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/operator/random/shuffle_op.cc b/src/operator/random/shuffle_op.cc
@@ -46,8 +46,13 @@ namespace {
 template<typename DType, typename Rand>
 void Shuffle1D(DType* const out, const index_t size, Rand* const prnd) {
   #ifdef USE_GNU_PARALLEL_SHUFFLE
-    auto rand_n = [prnd](index_t n) {
-      std::uniform_int_distribution<index_t> dist(0, n - 1);
+     /*
+      * See issue #15029: the data type of n needs to be compatible with
+      * the gcc library: https://github.com/gcc-mirror/gcc/blob/master/libstdc%2B%2B\
+      * -v3/include/parallel/random_shuffle.h#L384
+      */
+    auto rand_n = [prnd](uint32_t n) {
+      std::uniform_int_distribution<uint32_t> dist(0, n - 1);
       return dist(*prnd);
     };
     __gnu_parallel::random_shuffle(out, out + size, rand_n);

diff --git a/tests/python/unittest/test_random.py b/tests/python/unittest/test_random.py
@@ -867,6 +867,8 @@ def testLarge(data, repeat):
     # Test larger arrays
     testLarge(mx.nd.arange(0, 100000).reshape((10, 10000)), 10)
     testLarge(mx.nd.arange(0, 100000).reshape((10000, 10)), 10)
+    testLarge(mx.nd.arange(0, 100000), 10)
+
 
 @with_seed()
 def test_randint():