Large array support for randint (#14242)

* large array support for randint

* with seed for 2 random large array tests

* Trigger notification

src/operator/random/sampler.h

@@ -96,8 +96,8 @@ struct UniformSampler {
 template<typename xpu>
 struct SampleRandIntKernel {
   template<typename IType, typename OType>
-  MSHADOW_XINLINE static void Map(int id, RandGenerator<xpu, OType> gen,
-                                  const int N, const int step,
+  MSHADOW_XINLINE static void Map(index_t id, RandGenerator<xpu, OType> gen,
+                                  const index_t N, const index_t step,
                                   index_t nParm, index_t nSample,
                                   const IType *lower, const IType *upper, OType *out) {
     RNG_KERNEL_LOOP(xpu, OType, id, gen, N, step, {

tests/nightly/test_large_array.py

@@ -18,6 +18,7 @@
 import mxnet as mx
 import numpy as np
 from mxnet import gluon, nd
+from tests.python.unittest.common import with_seed
 
 # dimension constants
 MEDIUM_X = 10000
@@ -45,10 +46,23 @@ def test_ndarray_ones():
     assert a[-1][0] == 1
     assert nd.sum(a).asnumpy() == LARGE_SIZE
 
+@with_seed()
 def test_ndarray_random_uniform():
     a = nd.random.uniform(shape=(LARGE_X, SMALL_Y))
     assert a[-1][0] != 0
 
+@with_seed()
+def test_ndarray_random_randint():
+    a = nd.random.randint(100, 10000, shape=(LARGE_X, SMALL_Y))
+    assert a.shape == (LARGE_X, SMALL_Y)
+    # check if randint can generate value greater than 2**32 (large)
+    low_large_value = 2**32
+    high_large_value = 2**34
+    a = nd.random.randint(low_large_value,high_large_value)
+    low = mx.nd.array([low_large_value],dtype='int64')
+    high = mx.nd.array([high_large_value],dtype='int64')
+    assert a.__gt__(low) & a.__lt__(high)
+
 def test_ndarray_empty():
     a = nd.empty((LARGE_X, SMALL_Y))
     assert a.shape == (LARGE_X, SMALL_Y)