Repository: cupy/cupy

Issue:
random.randint support for range larger than int32
[random.randint](https://github.com/cupy/cupy/blob/master/cupy/random/generator.py#L1122) doesn't support ranges larger than int32 apparently because [RandomState._interval](https://github.com/cupy/cupy/blob/f121d14f6d3d2f5b25a58f26d2f8bca40dd10812/cupy/random/generator.py#L631) doesn't. Is there a plan to support larger numbers such as int64 and uint64?


Discussion:
Hello and thanks for your question,
The main problem here is the reliance in the `curand` library which does not support generating numbers with a extent greater than 2^32 bits with the default generator. We have to rely on a different one.
I am working on a solution to this.

Resolving patch:
diff --git a/cupy/random/generator.py b/cupy/random/generator.py
--- a/cupy/random/generator.py
+++ b/cupy/random/generator.py
@@ -604,11 +604,7 @@ def _interval(self, mx, size):
             If ``int``, 1-D array of length size is returned.
             If ``tuple``, multi-dimensional array with shape
             ``size`` is returned.
-            Currently, only 32 bit integers can be sampled.
-            If 0 :math:`\\leq` ``mx`` :math:`\\leq` 0x7fffffff,
-            a ``numpy.int32`` array is returned.
-            If 0x80000000 :math:`\\leq` ``mx`` :math:`\\leq` 0xffffffff,
-            a ``numpy.uint32`` array is returned.
+            Currently, only 32 bit or 64 bit integers can be sampled.
         """  # NOQA
         if size is None:
             return self._interval(mx, 1).reshape(())
@@ -618,18 +614,18 @@ def _interval(self, mx, size):
             size = (size, )
 
         if mx == 0:
-            return cupy.zeros(size, dtype=numpy.int32)
+            return cupy.zeros(size, dtype=numpy.uint32)
 
         if mx < 0:
             raise ValueError(
                 'mx must be non-negative (actual: {})'.format(mx))
-        elif mx <= 0x7fffffff:
-            dtype = numpy.int32
         elif mx <= 0xffffffff:
             dtype = numpy.uint32
+        elif mx <= 0xffffffffffffffff:
+            dtype = numpy.uint64
         else:
             raise ValueError(
-                'mx must be within uint32 range (actual: {})'.format(mx))
+                'mx must be within uint64 range (actual: {})'.format(mx))
 
         mask = (1 << mx.bit_length()) - 1
         mask = cupy.array(mask, dtype=dtype)
@@ -640,11 +636,13 @@ def _interval(self, mx, size):
             return cupy.array(())
 
         sample = cupy.empty((n,), dtype=dtype)
+        size32 = sample.view(dtype=numpy.uint32).size
         n_rem = n  # The number of remaining elements to sample
         ret = None
         while n_rem > 0:
+            # Call 32-bit RNG to fill 32-bit or 64-bit `sample`
             curand.generate(
-                self._generator, sample.data.ptr, sample.size)
+                self._generator, sample.data.ptr, size32)
             # Drop the samples that exceed the upper limit
             sample &= mask
             success = sample <= mx
@@ -1118,10 +1116,6 @@ def randint(self, low, high=None, size=None, dtype='l'):
                 'high is out of bounds for {}'.format(cupy.dtype(dtype).name))
 
         diff = hi - lo - 1
-        if diff > cupy.iinfo(cupy.int32).max - cupy.iinfo(cupy.int32).min + 1:
-            raise NotImplementedError(
-                'Sampling from a range whose extent is larger than int32 '
-                'range is currently not supported')
         x = self._interval(diff, size).astype(dtype, copy=False)
         cupy.add(x, lo, out=x)
         return x


Test patch:
diff --git a/tests/cupy_tests/random_tests/test_generator.py b/tests/cupy_tests/random_tests/test_generator.py
--- a/tests/cupy_tests/random_tests/test_generator.py
+++ b/tests/cupy_tests/random_tests/test_generator.py
@@ -863,7 +863,7 @@ def test_zero(self):
         vals = self.generate_many(0, shape, _count=10)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == shape
             assert (val == 0).all()
 
@@ -872,7 +872,7 @@ def test_shape_zero(self):
         vals = self.generate_many(mx, None, _count=10)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == ()
             assert (0 <= val).all()
             assert (val <= mx).all()
@@ -884,7 +884,7 @@ def test_shape_one_dim(self):
         vals = self.generate_many(mx, size, _count=10)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == (size,)
             assert (0 <= val).all()
             assert (val <= mx).all()
@@ -896,31 +896,17 @@ def test_shape_multi_dim(self):
         vals = self.generate_many(mx, shape, _count=10)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == shape
             assert (0 <= val).all()
             assert (val <= mx).all()
         # TODO(niboshi): Distribution test
 
-    def test_int32_range(self):
-        v = self.generate(0x00000000, 2)
-        assert v.dtype == numpy.int32
-
-        v = self.generate(0x7fffffff, 2)
-        assert v.dtype == numpy.int32
-
-    def test_uint32_range(self):
-        v = self.generate(0x80000000, 2)
-        assert v.dtype == numpy.uint32
-
-        v = self.generate(0xffffffff, 2)
-        assert v.dtype == numpy.uint32
-
     def test_bound_1(self):
         vals = self.generate_many(10, (2, 3), _count=10)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == (2, 3)
             assert (0 <= val).all()
             assert (val <= 10).all()
@@ -929,7 +915,7 @@ def test_bound_2(self):
         vals = self.generate_many(2, None, _count=20)
         for val in vals:
             assert isinstance(val, cupy.ndarray)
-            assert val.dtype == numpy.int32
+            assert val.dtype.kind in 'iu'
             assert val.shape == ()
             assert (0 <= val).all()
             assert (val <= 2).all()
@@ -1191,6 +1177,9 @@ def test_randint_2(self):
     def test_randint_3(self):
         self.generate(3, 10, size=0)
 
+    def test_randint_int64_1(self):
+        self.generate(2**34, 2**40, 3)
+
 
 @testing.gpu
 @testing.fix_random()