Repository: cupy/cupy

Issue:
Allow RandomState.permutation() to accept cupy.ndarray as input
Can you please allow RandomState.permutation() accept cupy.ndarray as input

* Code to reproduce
import cupy as cp
rng = cp.random.generator.RandomState(seed=0)
rng.permutation(cp.array([1, 2, 3]))

* Error messages, stack traces, or logs
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/rapids/notebooks/k.py", line 1174, in split
    for train, test in self._iter_indices(X, y, groups):
  File "/rapids/notebooks/k.py", line 1494, in _iter_indices
    print(rng.permutation(cp.array([1, 2, 3])))
  File "/conda/envs/rapids/lib/python3.6/site-packages/cupy/random/generator.py", line 956, in permutation
    raise TypeError('The data type of argument "num" must be integer')
TypeError: The data type of argument "num" must be integer



Discussion:

Resolving patch:
diff --git a/cupy/random/generator.py b/cupy/random/generator.py
--- a/cupy/random/generator.py
+++ b/cupy/random/generator.py
@@ -948,13 +948,17 @@ def shuffle(self, a):
         if a.ndim == 0:
             raise TypeError('An array whose ndim is 0 is not supported')
 
-        a[:] = a[self.permutation(len(a))]
+        a[:] = a[self._permutation(len(a))]
 
-    def permutation(self, num):
-        """Returns a permuted range."""
-        if not isinstance(num, six.integer_types):
-            raise TypeError('The data type of argument "num" must be integer')
+    def permutation(self, a):
+        """Returns a permuted range or a permutation of an array."""
+        if isinstance(a, six.integer_types):
+            return self._permutation(a)
+        else:
+            return a[self._permutation(len(a))]
 
+    def _permutation(self, num):
+        """Returns a permuted range."""
         sample = cupy.empty((num), dtype=numpy.int32)
         curand.generate(self._generator, sample.data.ptr, num)
         if 128 < num <= 32 * 1024 * 1024:
diff --git a/cupy/random/permutations.py b/cupy/random/permutations.py
--- a/cupy/random/permutations.py
+++ b/cupy/random/permutations.py
@@ -1,5 +1,4 @@
 from cupy.random import generator
-import six
 
 
 def shuffle(a):
@@ -29,7 +28,4 @@ def permutation(a):
     .. seealso:: :func:`numpy.random.permutation`
     """
     rs = generator.get_random_state()
-    if isinstance(a, six.integer_types):
-        return rs.permutation(a)
-    else:
-        return a[rs.permutation(len(a))]
+    return rs.permutation(a)


Test patch:
diff --git a/tests/cupy_tests/random_tests/test_permutations.py b/tests/cupy_tests/random_tests/test_permutations.py
--- a/tests/cupy_tests/random_tests/test_permutations.py
+++ b/tests/cupy_tests/random_tests/test_permutations.py
@@ -7,31 +7,44 @@
 import numpy
 
 
+@testing.parameterize(
+    {'seed': None},
+    {'seed': 0},
+)
 @testing.gpu
 class TestPermutations(unittest.TestCase):
 
+    def _xp_random(self, xp):
+        if self.seed is None:
+            return xp.random
+        else:
+            return xp.random.RandomState(seed=self.seed)
+
     # Test ranks
 
     @testing.numpy_cupy_raises()
     def test_permutation_zero_dim(self, xp):
+        xp_random = self._xp_random(xp)
         a = testing.shaped_random((), xp)
-        xp.random.permutation(a)
+        xp_random.permutation(a)
 
     # Test same values
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     def test_permutation_sort_1dim(self, dtype):
+        cupy_random = self._xp_random(cupy)
         a = cupy.arange(10, dtype=dtype)
         b = cupy.copy(a)
-        c = cupy.random.permutation(a)
+        c = cupy_random.permutation(a)
         testing.assert_allclose(a, b)
         testing.assert_allclose(b, cupy.sort(c))
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     def test_permutation_sort_ndim(self, dtype):
+        cupy_random = self._xp_random(cupy)
         a = cupy.arange(15, dtype=dtype).reshape(5, 3)
         b = cupy.copy(a)
-        c = cupy.random.permutation(a)
+        c = cupy_random.permutation(a)
         testing.assert_allclose(a, b)
         testing.assert_allclose(b, cupy.sort(c, axis=0))
 
@@ -41,10 +54,16 @@ def test_permutation_sort_ndim(self, dtype):
     def test_permutation_seed1(self, dtype):
         a = testing.shaped_random((10,), cupy, dtype)
         b = cupy.copy(a)
-        cupy.random.seed(0)
-        pa = cupy.random.permutation(a)
-        cupy.random.seed(0)
-        pb = cupy.random.permutation(b)
+
+        cupy_random = self._xp_random(cupy)
+        if self.seed is None:
+            cupy_random.seed(0)
+        pa = cupy_random.permutation(a)
+        cupy_random = self._xp_random(cupy)
+        if self.seed is None:
+            cupy_random.seed(0)
+        pb = cupy_random.permutation(b)
+
         testing.assert_allclose(pa, pb)