diff options
| -rw-r--r-- | numpy/random/generator.pyx | 6 | ||||
| -rw-r--r-- | numpy/random/tests/test_generator_mt19937.py | 9 |
2 files changed, 12 insertions, 3 deletions
diff --git a/numpy/random/generator.pyx b/numpy/random/generator.pyx index 37ac57c06..df7485a97 100644 --- a/numpy/random/generator.pyx +++ b/numpy/random/generator.pyx @@ -3786,7 +3786,7 @@ cdef class Generator: # Shuffling and permutations: def shuffle(self, object x, axis=0): """ - shuffle(x) + shuffle(x, axis=0) Modify a sequence in-place by shuffling its contents. @@ -3858,7 +3858,7 @@ cdef class Generator: x = np.swapaxes(x, 0, axis) buf = np.empty_like(x[0, ...]) with self.lock: - for i in reversed(range(1, n)): + for i in reversed(range(1, len(x))): j = random_interval(&self._bitgen, i) if i == j: # i == j is not needed and memcpy is undefined. @@ -3928,7 +3928,7 @@ cdef class Generator: def permutation(self, object x, axis=0): """ - permutation(x) + permutation(x, axis=0) Randomly permute a sequence, or return a permuted range. diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py index 20bc10cd0..391c33c1a 100644 --- a/numpy/random/tests/test_generator_mt19937.py +++ b/numpy/random/tests/test_generator_mt19937.py @@ -746,6 +746,15 @@ class TestRandomDist(object): random.shuffle(actual, axis=-1) assert_array_equal(actual, desired) + def test_shuffle_axis_nonsquare(self): + y1 = np.arange(20).reshape(2, 10) + y2 = y1.copy() + random = Generator(MT19937(self.seed)) + random.shuffle(y1, axis=1) + random = Generator(MT19937(self.seed)) + random.shuffle(y2.T) + assert_array_equal(y1, y2) + def test_shuffle_masked(self): # gh-3263 a = np.ma.masked_values(np.reshape(range(20), (5, 4)) % 3 - 1, -1) |