Repository: cupy/cupy

Issue:
Feature request: function to set a random state
You can get the current random generator state of cupy via `cupy.random.get_random_state`, but there seems no function available to set a random state to cupy.

numpy has both `numpy.random.get_state` and `numpy.random.set_state`.

The implementation below works for now, but it depends on undocumented things and I think it's better to have an official function for it.
```
def cupy_set_random_state(random_state):
    device_id = cupy.cuda.Device().id
    cupy.random.generator._random_states[device_id] = random_state
```


Discussion:

Resolving patch:
diff --git a/cupy/random/__init__.py b/cupy/random/__init__.py
--- a/cupy/random/__init__.py
+++ b/cupy/random/__init__.py
@@ -25,6 +25,7 @@ def bytes(length):
 from cupy.random.generator import RandomState  # NOQA
 from cupy.random.generator import reset_states  # NOQA
 from cupy.random.generator import seed  # NOQA
+from cupy.random.generator import set_random_state  # NOQA
 from cupy.random.permutations import shuffle  # NOQA
 from cupy.random.sample import choice  # NOQA
 from cupy.random.sample import multinomial  # NOQA
diff --git a/cupy/random/generator.py b/cupy/random/generator.py
--- a/cupy/random/generator.py
+++ b/cupy/random/generator.py
@@ -12,6 +12,7 @@
 from cupy import core
 from cupy import cuda
 from cupy.cuda import curand
+from cupy.cuda import device
 
 
 class RandomState(object):
@@ -439,6 +440,19 @@ def get_random_state():
     return rs
 
 
+def set_random_state(rs):
+    """Sets the state of the random number generator for the current device.
+
+    Args:
+        state(RandomState): Random state to set for the current device.
+    """
+    if not isinstance(rs, RandomState):
+        raise TypeError(
+            'Random state must be an instance of RandomState. '
+            'Actual: {}'.format(type(rs)))
+    _random_states[device.get_device_id()] = rs
+
+
 def _check_and_get_dtype(dtype):
     dtype = numpy.dtype(dtype)
     if dtype.char not in ('f', 'd'):


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
@@ -524,12 +524,33 @@ def test_get_random_state_memoized(self):
 
 
 @testing.gpu
-class TestGetRandomStateThreadSafe(unittest.TestCase):
+class TestSetRandomState(unittest.TestCase):
+
+    def setUp(self):
+        self.rs_tmp = generator._random_states
+
+    def tearDown(self, *args):
+        generator._random_states = self.rs_tmp
+
+    def test_set_random_state(self):
+        rs = generator.RandomState()
+        generator.set_random_state(rs)
+        assert generator.get_random_state() is rs
+
+    def test_set_random_state_call_multiple_times(self):
+        generator.set_random_state(generator.RandomState())
+        rs = generator.RandomState()
+        generator.set_random_state(rs)
+        assert generator.get_random_state() is rs
+
+
+@testing.gpu
+class TestRandomStateThreadSafe(unittest.TestCase):
 
     def setUp(self):
         cupy.random.reset_states()
 
-    def test_thread_safe(self):
+    def test_get_random_state_thread_safe(self):
         seed = 10
         threads = [
             threading.Thread(target=lambda: cupy.random.seed(seed)),
@@ -551,6 +572,20 @@ def test_thread_safe(self):
         expected = cupy.random.uniform()
         self.assertEqual(actual, expected)
 
+    def test_set_random_state_thread_safe(self):
+        rs = cupy.random.RandomState()
+        threads = [
+            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
+            threading.Thread(target=lambda: cupy.random.set_random_state(rs)),
+        ]
+
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert cupy.random.get_random_state() is rs
+
 
 @testing.gpu
 class TestGetRandomState2(unittest.TestCase):