Advance rng state in permutation (dask#12031)

Author: jrbourbeau

dask/array/random.py

@@ -791,9 +791,6 @@ def _rng_from_bitgen(bitgen):
 
 
 def _shuffle(bit_generator, x, axis=0):
-    state_data = bit_generator.state
-    bit_generator = type(bit_generator)()
-    bit_generator.state = state_data
     state = _rng_from_bitgen(bit_generator)
     return state.shuffle(x, axis=axis)
 

dask/array/tests/test_random.py

@@ -403,6 +403,14 @@ def test_permutation(generator_class):
     r2 = b.permutation(x)
     assert_eq(r1, r2)
 
+    # Ensure subsequent `permutation` calls return a different result
+    # See https://github.com/dask/dask/issues/12029
+    r3 = a.permutation(x)
+    r4 = b.permutation(x)
+    assert_eq(r3, r4)
+    with pytest.raises(AssertionError):
+        assert_eq(r1, r3)
+
     x = generator_class().permutation(100)
     assert x.shape == (100,)