sklam · Sep 21, 2022 · Aug 8, 2022 · Aug 9, 2022 · Aug 22, 2022 · Aug 29, 2022
diff --git a/numba/np/random/generator_methods.py b/numba/np/random/generator_methods.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 from numba.core import types
-from numba.core.extending import overload_method
+from numba.core.extending import overload_method, register_jitable
 from numba.np.numpy_support import as_dtype, from_dtype
 from numba.np.random.generator_core import next_float, next_double
 from numba.np.numpy_support import is_nonelike
@@ -22,6 +22,7 @@
      random_geometric, random_zipf, random_triangular,
      random_poisson, random_negative_binomial)
 from numba.np.random import random_methods
+from numba.cpython.unsafe.tuple import tuple_setitem
 
 
 def _get_proper_func(func_32, func_64, dtype, dist_name="the given"):
@@ -166,6 +167,54 @@ def impl(inst, low, high, size=None,
         return impl
 
 
+# Overload the Generator().shuffle()
+@overload_method(types.NumPyRandomGeneratorType, 'shuffle')
+def NumPyRandomGeneratorType_shuffle(inst, arr, axis=0):
+    check_types(arr, [types.Array], 'arr')
+    check_types(axis, [int, types.Integer], 'axis')
+
+    axis_impl = tuple([slice(None, None)] * arr.ndim)
+
+    def impl(inst, arr, axis=0):
+
+        for i in range(arr.shape[axis] - 1, 0, -1):
+            j = types.intp(random_methods.random_interval(inst.bit_generator,
+                                                          i))
+            if i == j:
+                continue
+            axis_impl1 = tuple_setitem(axis_impl, axis, slice(i, i + 1))
+            axis_impl2 = tuple_setitem(axis_impl, axis, slice(j, j + 1))
+
+            buffer = np.copy(arr[axis_impl2])
+            arr[axis_impl2] = np.copy(arr[axis_impl1])
+            arr[axis_impl1] = buffer
+
+    return impl
+
+
+# Overload the Generator().permutation()
+@overload_method(types.NumPyRandomGeneratorType, 'permutation')
+def NumPyRandomGeneratorType_permutation(inst, arr, axis=0):
+    check_types(arr, [types.Array], 'arr')
+    check_types(axis, [int, types.Integer], 'axis')
+
+    if isinstance(arr, types.Integer):
+        @register_jitable
+        def array_maker(arr):
+            return np.arange(arr)
+    else:
+        @register_jitable
+        def array_maker(arr):
+            return arr.copy()
+
+    def impl(inst, arr, axis=0):
+        new_arr = array_maker(arr)
+        inst.shuffle(new_arr, axis=axis)
+        return new_arr
+
+    return impl
+
+
 # Overload the Generator().random()
 @overload_method(types.NumPyRandomGeneratorType, 'random')
 def NumPyRandomGeneratorType_random(inst, size=None, dtype=np.float64):

diff --git a/numba/np/random/random_methods.py b/numba/np/random/random_methods.py
@@ -11,6 +11,18 @@
 # https://github.com/numpy/numpy/blob/7cfef93c77599bd387ecc6a15d186c5a46024dac/numpy/random/src/distributions/distributions.c
 
 
+@register_jitable
+def gen_mask(max):
+    mask = uint64(max)
+    mask |= mask >> 1
+    mask |= mask >> 2
+    mask |= mask >> 4
+    mask |= mask >> 8
+    mask |= mask >> 16
+    mask |= mask >> 32
+    return mask
+
+
 @register_jitable
 def buffered_bounded_bool(bitgen, off, rng, bcnt, buf):
     if (rng == 0):
@@ -331,3 +343,30 @@ def _randint_arg_check(low, high, endpoint, lower_bound, upper_bound):
             raise ValueError("high is out of bounds")
         if low > high:  # -1 already subtracted, closed interval
             raise ValueError("low is greater than high in given interval")
+
+
+@register_jitable
+def random_interval(bitgen, max_val):
+    if (max_val == 0):
+        return 0
+
+    max_val = uint64(max_val)
+    mask = uint64(max_val)
+
+    mask |= mask >> 1
+    mask |= mask >> 2
+    mask |= mask >> 4
+    mask |= mask >> 8
+    mask |= mask >> 16
+    mask |= mask >> 32
+
+    if (max_val <= 0xffffffff):
+        value = uint64(next_uint32(bitgen)) & mask
+        while value > max_val:
+            value = uint64(next_uint32(bitgen)) & mask
+    else:
+        value = next_uint64(bitgen) & mask
+        while value > max_val:
+            value = next_uint64(bitgen) & mask
+
+    return uint64(value)
diff --git a/numba/tests/test_np_randomgen.py b/numba/tests/test_np_randomgen.py
@@ -946,6 +946,50 @@ def test_negative_binomial(self):
         self._check_invalid_types(dist_func, ['n', 'p', 'size'],
                                   [1, 0.75, (1,)], ['x', 'x', ('x',)])
 
+    def test_shuffle(self):
+        test_sizes = [(10, 20, 30)]
+        bitgen_types = [None, MT19937]
+
+        def dist_func(x, size, dtype):
+            arr = x.random(size=size)
+            x.shuffle(arr, axis=1)
+            return arr
+
+        for _size in test_sizes:
+            for _bitgen in bitgen_types:
+                with self.subTest(_size=_size, _bitgen=_bitgen):
+                    self.check_numpy_parity(dist_func, _bitgen,
+                                            None, _size, None,
+                                            adjusted_ulp_prec)
+
+        def dist_func(x, arr, axis):
+            x.shuffle(arr, axis=axis)
+            return arr
+
+        self._check_invalid_types(dist_func, ['arr', 'axis'],
+                                  [np.array([3,4,5]), 0], ['x', 'x'])
+
+    def test_permutation(self):
+        test_sizes = [(10, 20, 30)]
+        bitgen_types = [None, MT19937]
+
+        def dist_func(x, size, dtype):
+            arr = x.random(size=size)
+            return x.permutation(arr, axis=1)
+
+        for _size in test_sizes:
+            for _bitgen in bitgen_types:
+                with self.subTest(_size=_size, _bitgen=_bitgen):
+                    self.check_numpy_parity(dist_func, _bitgen,
+                                            None, _size, None,
+                                            adjusted_ulp_prec)
+
+        def dist_func(x, arr, axis):
+            return x.permutation(arr, axis=axis)
+
+        self._check_invalid_types(dist_func, ['arr', 'axis'],
+                                  [np.array([3,4,5]), 0], ['x', 'x'])
+
 
 class TestGeneratorCaching(TestCase, SerialMixin):
     def test_randomgen_caching(self):