BUG: Make np.partition and np.sort work on np.matrix when axis=None

Both were making the normally valid assumption that flatten actually flattens, which turns out to be false for matrices.

Old behavior:

>>> a = np.matrix([[1, 2, 0]])
>>> np.partition(a, 1, axis=None)
ValueError: kth(=1) out of bounds (1)
>>> np.sort(a, axis=None)
matrix([[1, 2, 0]])

2 files changed, 19 insertions, 2 deletions

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index 43584349f..4dfeb35ca 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -657,8 +657,9 @@ def partition(a, kth, axis=-1, kind='introselect', order=None):
 
     """
     if axis is None:
+        # flatten returns (1, N) for np.matrix, so always use the last axis
         a = asanyarray(a).flatten()
-        axis = 0
+        axis = -1
     else:
         a = asanyarray(a).copy(order="K")
     a.partition(kth, axis=axis, kind=kind, order=order)
@@ -840,8 +841,9 @@ def sort(a, axis=-1, kind='quicksort', order=None):
 
     """
     if axis is None:
+        # flatten returns (1, N) for np.matrix, so always use the last axis
         a = asanyarray(a).flatten()
-        axis = 0
+        axis = -1
     else:
         a = asanyarray(a).copy(order="K")
     a.sort(axis=axis, kind=kind, order=order)
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 661561ab3..b3515ae6e 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1742,6 +1742,13 @@ class TestMethods(object):
         assert_equal(r, np.array([('a', 1), ('c', 3), ('b', 255), ('d', 258)],
                                  dtype=mydtype))
 
+    def test_sort_matrix_none(self):
+        a = np.matrix([[2, 1, 0]])
+        actual = np.sort(a, axis=None)
+        expected = np.matrix([[0, 1, 2]])
+        assert_equal(actual, expected)
+        assert_(type(expected) is np.matrix)
+
     def test_argsort(self):
         # all c scalar argsorts use the same code with different types
         # so it suffices to run a quick check with one type. The number
@@ -2487,6 +2494,14 @@ class TestMethods(object):
                 assert_array_equal(np.partition(d, kth)[kth], tgt,
                                    err_msg="data: %r\n kth: %r" % (d, kth))
 
+    def test_partition_matrix_none(self):
+        # gh-4301
+        a = np.matrix([[2, 1, 0]])
+        actual = np.partition(a, 1, axis=None)
+        expected = np.matrix([[0, 1, 2]])
+        assert_equal(actual, expected)
+        assert_(type(expected) is np.matrix)
+
     def test_argpartition_gh5524(self):
         #  A test for functionality of argpartition on lists.
         d = [6,7,3,2,9,0]