sklam · Sep 21, 2022 · Aug 8, 2022 · Aug 9, 2022 · Aug 22, 2022 · Aug 29, 2022
diff --git a/docs/source/reference/numpysupported.rst b/docs/source/reference/numpysupported.rst
@@ -628,10 +628,12 @@ The following :py:class:`Generator` methods are supported:
 * :func:`numpy.random.Generator().negative_binomial()` (*)
 * :func:`numpy.random.Generator().normal()` (*)
 * :func:`numpy.random.Generator().pareto()`
+* :func:`numpy.random.Generator().permutation()` (Only accepts NumPy ndarrays and integers.)
 * :func:`numpy.random.Generator().poisson()` (*)
 * :func:`numpy.random.Generator().power()`
 * :func:`numpy.random.Generator().random()`
 * :func:`numpy.random.Generator().rayleigh()`
+* :func:`numpy.random.Generator().shuffle()` (Only accepts NumPy ndarrays.)
 * :func:`numpy.random.Generator().standard_cauchy()`
 * :func:`numpy.random.Generator().standard_exponential()`
 * :func:`numpy.random.Generator().standard_gamma()` (*)

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,59 @@ def impl(inst, low, high, size=None,
         return impl
 
 
+# The following `shuffle` implementation is a direct translation from:
+# https://github.com/numpy/numpy/blob/95e3e7f445407e4f355b23d6a9991d8774f0eb0c/numpy/random/_generator.pyx#L4578
+
+# Overload the Generator().shuffle()
+@overload_method(types.NumPyRandomGeneratorType, 'shuffle')
+def NumPyRandomGeneratorType_shuffle(inst, x, axis=0):
+    check_types(x, [types.Array], 'x')
+    check_types(axis, [int, types.Integer], 'axis')
+    axis_impl = tuple([slice(None, None)] * x.ndim)
+
+    def impl(inst, x, axis=0):
+        if axis > x.ndim - 1:
+            raise IndexError("Axis is out of bounds for the given array")
+        for i in range(x.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(x[axis_impl2])
+            x[axis_impl2] = np.copy(x[axis_impl1])
+            x[axis_impl1] = buffer
+
+    return impl
+
+
+# The following `permutation` implementation is a direct translation from:
+# https://github.com/numpy/numpy/blob/95e3e7f445407e4f355b23d6a9991d8774f0eb0c/numpy/random/_generator.pyx#L4710
+# Overload the Generator().permutation()
+@overload_method(types.NumPyRandomGeneratorType, 'permutation')
+def NumPyRandomGeneratorType_permutation(inst, x, axis=0):
+    check_types(x, [types.Array, types.Integer], 'x')
+    check_types(axis, [int, types.Integer], 'axis')
+
+    if isinstance(x, types.Integer):
+        @register_jitable
+        def array_maker(x):
+            return np.arange(x)
+    else:
+        @register_jitable
+        def array_maker(x):
+            return x.copy()
+
+    def impl(inst, x, axis=0):
+        new_arr = array_maker(x)
+        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
@@ -2,6 +2,7 @@
 import numpy as np
 import sys
 import platform
+import itertools
 
 from numba import types
 from numba.core.config import IS_32BITS
@@ -946,6 +947,105 @@ def test_negative_binomial(self):
         self._check_invalid_types(dist_func, ['n', 'p', 'size'],
                                   [1, 0.75, (1,)], ['x', 'x', ('x',)])
 
+    # NumPy tests at:
+    # https://github.com/numpy/numpy/blob/95e3e7f445407e4f355b23d6a9991d8774f0eb0c/numpy/random/tests/test_generator_mt19937.py#L936
+    # Written in following format for semblance with existing Generator tests.
+    def test_shuffle(self):
+        test_sizes = [(10, 20, 30)]
+        bitgen_types = [None, MT19937]
+        axes = [0, 1, 2]
+
+        for _size, _bitgen, _axis in itertools.product(test_sizes,
+                                                       bitgen_types,
+                                                       axes):
+            with self.subTest(_size=_size, _bitgen=_bitgen, _axis=_axis):
+                def dist_func(x, size, dtype):
+                    arr = x.random(size=size)
+                    x.shuffle(arr, axis=_axis)
+                    return arr
+                self.check_numpy_parity(dist_func, _bitgen,
+                                        None, _size, None,
+                                        0)
+
+    def test_shuffle_empty(self):
+        a = np.array([])
+        b = np.array([])
+
+        def dist_func(x, arr):
+            x.shuffle(arr)
+            return arr
+
+        nb_func = numba.njit(dist_func)
+        rng = lambda: np.random.default_rng(1)
+
+        np.testing.assert_equal(dist_func(rng(), a), nb_func(rng(), b))
+
+    def test_shuffle_check(self):
+        self.disable_leak_check()
+
+        def dist_func(x, arr, axis):
+            x.shuffle(arr, axis=axis)
+            return arr
+
+        self._check_invalid_types(dist_func, ['x', 'axis'],
+                                  [np.array([3,4,5]), 0], ['x', 'x'])
+
+        rng = np.random.default_rng(1)
+        with self.assertRaises(IndexError) as raises:
+            numba.njit(dist_func)(rng, np.array([3,4,5]), 2)
+        self.assertIn(
+            'Axis is out of bounds for the given array',
+            str(raises.exception)
+        )
+
+    # NumPy tests at:
+    # https://github.com/numpy/numpy/blob/95e3e7f445407e4f355b23d6a9991d8774f0eb0c/numpy/random/tests/test_generator_mt19937.py#L1030
+    # Written in following format for semblance with existing Generator tests.
+    def test_permutation(self):
+        test_sizes = [(10, 20, 30)]
+        bitgen_types = [None, MT19937]
+        axes = [0, 1, 2]
+
+        for _size, _bitgen, _axis in itertools.product(test_sizes,
+                                                       bitgen_types,
+                                                       axes):
+            with self.subTest(_size=_size, _bitgen=_bitgen, _axis=_axis):
+                def dist_func(x, size, dtype):
+                    arr = x.random(size=size)
+                    return x.permutation(arr, axis=1)
+                self.check_numpy_parity(dist_func, _bitgen,
+                                        None, _size, None,
+                                        0)
+
+    def test_permutation_exception(self):
+        self.disable_leak_check()
+
+        def dist_func(x, arr, axis):
+            return x.permutation(arr, axis=axis)
+
+        self._check_invalid_types(dist_func, ['x', 'axis'],
+                                  [np.array([3,4,5]), 0], ['x', 'x'])
+
+        rng = np.random.default_rng(1)
+        with self.assertRaises(IndexError) as raises:
+            numba.njit(dist_func)(rng, np.array([3,4,5]), 2)
+        self.assertIn(
+            'Axis is out of bounds for the given array',
+            str(raises.exception)
+        )
+
+    def test_permutation_empty(self):
+        a = np.array([])
+        b = np.array([])
+
+        def dist_func(x, arr):
+            return x.permutation(arr)
+
+        nb_func = numba.njit(dist_func)
+        rng = lambda: np.random.default_rng(1)
+
+        np.testing.assert_equal(dist_func(rng(), a), nb_func(rng(), b))
+
 
 class TestGeneratorCaching(TestCase, SerialMixin):
     def test_randomgen_caching(self):