BUG: random: Use correct length when axis is given to shuffle.

When an axis argument was given, shuffle was using the original length of
the array instead of the length of the given axis.  This meant that, for
example, if an array with shape (2, 10) was shuffled with axis=1, only the
first two columns were shuffled.

1 files changed, 9 insertions, 0 deletions

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)