Repository: numba/numba

Issue:
Support np.random.permutation
`np.random.permutation` is useful for shuffling training data of deep learning workloads.
@agayev


Discussion:

Resolving patch:
diff --git a/numba/targets/randomimpl.py b/numba/targets/randomimpl.py
--- a/numba/targets/randomimpl.py
+++ b/numba/targets/randomimpl.py
@@ -1145,26 +1145,56 @@ def zipf_impl(a):
     res = context.compile_internal(builder, zipf_impl, sig, args)
     return impl_ret_untracked(context, builder, sig.return_type, res)
 
-
-@lower("random.shuffle", types.Buffer)
-def shuffle_impl(context, builder, sig, args):
-    res = _shuffle_impl(context, builder, sig, args, random.randrange)
-    return impl_ret_untracked(context, builder, sig.return_type, res)
-
-@lower("np.random.shuffle", types.Buffer)
-def shuffle_impl(context, builder, sig, args):
-    res = _shuffle_impl(context, builder, sig, args, np.random.randint)
-    return impl_ret_untracked(context, builder, sig.return_type, res)
-
-def _shuffle_impl(context, builder, sig, args, _randrange):
-    def shuffle_impl(arr):
-        i = arr.shape[0] - 1
-        while i > 0:
-            j = _randrange(i + 1)
-            arr[i], arr[j] = arr[j], arr[i]
-            i -= 1
-
-    return context.compile_internal(builder, shuffle_impl, sig, args)
+def do_shuffle_impl(arr, rng):
+
+    if not isinstance(arr, types.Buffer):
+        raise TypeError("The argument to shuffle() should be a buffer type")
+
+    if rng == "np":
+        rand = np.random.randint
+    elif rng == "py":
+        rand = random.randrange
+
+    if arr.ndim == 1:
+        def impl(arr):
+            i = arr.shape[0] - 1
+            while i > 0:
+                j = rand(i + 1)
+                arr[i], arr[j] = arr[j], arr[i]
+                i -= 1
+    else:
+        def impl(arr):
+            i = arr.shape[0] - 1
+            while i > 0:
+                j = rand(i + 1)
+                arr[i], arr[j] = np.copy(arr[j]), np.copy(arr[i])
+                i -= 1
+
+    return impl
+
+@overload(random.shuffle)
+def shuffle_impl(arr):
+    return do_shuffle_impl(arr, "py")
+
+@overload(np.random.shuffle)
+def shuffle_impl(arr):
+    return do_shuffle_impl(arr, "np")
+
+@overload(np.random.permutation)
+def permutation_impl(arr):
+    if isinstance(arr, types.Integer):
+        def permutation_impl(n):
+            arr = np.arange(n)
+            np.random.shuffle(arr)
+            return arr
+    elif isinstance(arr, types.Array):
+        def permutation_impl(arr):
+            arr_copy = arr.copy()
+            np.random.shuffle(arr_copy)
+            return arr_copy
+    else:
+        permutation_impl = None
+    return permutation_impl
 
 
 # ------------------------------------------------------------------------
diff --git a/numba/typing/randomdecl.py b/numba/typing/randomdecl.py
--- a/numba/typing/randomdecl.py
+++ b/numba/typing/randomdecl.py
@@ -284,16 +284,5 @@ class Random_triangular(ConcreteTemplate):
     cases = [signature(tp, tp, tp) for tp in _float_types]
     cases += [signature(tp, tp, tp, tp) for tp in _float_types]
 
-# Other
-
-@infer_global(random.shuffle, typing_key="random.shuffle")
-@infer_global(np.random.shuffle, typing_key="np.random.shuffle")
-class Random_shuffle(AbstractTemplate):
-    def generic(self, args, kws):
-        arr, = args
-        if isinstance(arr, types.Buffer) and arr.ndim == 1 and arr.mutable:
-            return signature(types.void, arr)
-
-
 # NOTE: some functions can have @overloads in numba.targets.randomimpl,
 # and therefore don't need a typing declaration here.


Test patch:
diff --git a/numba/tests/test_random.py b/numba/tests/test_random.py
--- a/numba/tests/test_random.py
+++ b/numba/tests/test_random.py
@@ -76,6 +76,21 @@ def numpy_check_randn(seed, a, b):
     got = np.random.randn(a, b)
     return expected, got
 
+# Numba's current PRNG for NumPy follows Python's implementation,
+# however, Python's implementation of shuffle does not work on
+# multidimensional arrays.  Hence, we cannot test np.random.shuffle with
+# Python's shuffle.  Therefore, we implement a shuffle in Python that
+# works on multidimensional arrays and compare against that.
+def python_shuffle(arr, r):
+    if arr.ndim == 1:
+        r.shuffle(arr)
+    else:
+        i = arr.shape[0] - 1
+        while i > 0:
+            j = r.randrange(i + 1)
+            arr[i], arr[j] = np.copy(arr[j]), np.copy(arr[i])
+            i -= 1
+    return arr
 
 def jit_with_args(name, argstring):
     code = """def func(%(argstring)s):
@@ -829,27 +844,31 @@ def test_numpy_zipf(self):
 
     def _check_shuffle(self, func, ptr):
         """
-        Check a shuffle()-like function for 1D arrays.
+        Check a shuffle()-like function for arrays.
         """
         # Our implementation follows Python 3's.
-        a = np.arange(20)
+        arrs = [np.arange(20), np.arange(32).reshape((8, 4))]
         if sys.version_info >= (3,):
             r = self._follow_cpython(ptr)
-            for i in range(3):
-                got = a.copy()
-                expected = a.copy()
-                func(got)
-                r.shuffle(expected)
-                self.assertPreciseEqual(got, expected)
+            for a in arrs:
+                for i in range(3):
+                    got = a.copy()
+                    expected = a.copy()
+                    func(got)
+                    python_shuffle(expected, r)
+                    self.assertPreciseEqual(got, expected)
         else:
             # Sanity check
-            for i in range(3):
-                b = a.copy()
-                func(b)
-                self.assertNotEqual(list(a), list(b))
-                self.assertEqual(sorted(a), sorted(b))
-                a = b
+            for a in arrs:
+                for i in range(3):
+                    b = a.copy()
+                    func(b)
+                    self.assertFalse(np.array_equal(a, b))
+                    self.assertTrue(np.array_equal(np.sort(a, axis=0),
+                                                   np.sort(b, axis=0)))
+                    a = b
         # Test with an arbitrary buffer-providing object
+        a = arrs[0]
         b = a.copy()
         func(memoryview(b))
         self.assertNotEqual(list(a), list(b))
@@ -898,6 +917,51 @@ def test_numpy_random_startup(self):
     def test_numpy_gauss_startup(self):
         self._check_startup_randomness("numpy_normal", (1.0, 1.0))
 
+    def test_numpy_random_permutation(self):
+
+        # Numba's current PRNG for NumPy follows Python's implementation, hence,
+        # we cannot test np.random.permutation against NumPy.  Since Python does
+        # not have a permutation implementation, with which we can test
+        # np.random.permutation, we implement one, just for testing purposes.
+        def python_permutation(r, x):
+            if isinstance(x, int):
+                arr = np.arange(x)
+            else:
+                arr = np.array(x)
+            python_shuffle(arr, r)
+            return arr
+
+        # Our implementation follows Python 3's.
+        func = jit_unary("np.random.permutation")
+        if sys.version_info >= (3,):
+            r = self._follow_cpython(get_np_state_ptr())
+            for s in [5, 10, 15, 20]:
+                a = np.arange(s)
+                b = a.copy()
+                # Test array version
+                self.assertPreciseEqual(func(a), python_permutation(r, a))
+                # Test int version
+                self.assertPreciseEqual(func(s), python_permutation(r, s))
+                # Permutation should not modify its argument
+                self.assertPreciseEqual(a, b)
+            # Check multi-dimensional arrays
+            arrs = [np.arange(10).reshape(2, 5),
+                    np.arange(27).reshape(3, 3, 3),
+                    np.arange(36).reshape(2, 3, 3, 2)]
+            for a in arrs:
+                b = a.copy()
+                self.assertPreciseEqual(func(a), python_permutation(r, a))
+                self.assertPreciseEqual(a, b)
+        else:
+            # Sanity check
+            arrs = [np.arange(20), np.arange(20).reshape(5, 2, 2)]
+            for a in arrs:
+                for i in range(3):
+                    b = func(a)
+                    self.assertFalse(np.array_equal(a, b))
+                    self.assertTrue(np.array_equal(np.sort(a, axis=0),
+                                                   np.sort(b, axis=0)))
+
 
 class TestRandomArrays(BaseTest):
     """