Fix crash in random.shuffle operator (#15041)

* fix crash in random_shuffle caused by int overflow * add unit test * add comment * remove small random test to avoid CI failure
apache · May 23, 2019 · 66aa983 · 66aa983
1 parent 503c750
commit 66aa983
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():