BUG: np.ma.trace gives the wrong result on ND arrays

Fixes #5560

3 files changed, 12 insertions, 7 deletions

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index 6f7c45859..a94be7b4d 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -1245,13 +1245,13 @@ def diagonal(a, offset=0, axis1=0, axis2=1):
     Returns
     -------
     array_of_diagonals : ndarray
-        If `a` is 2-D and not a matrix, a 1-D array of the same type as `a`
-        containing the diagonal is returned. If `a` is a matrix, a 1-D
+        If `a` is 2-D and not a `matrix`, a 1-D array of the same type as `a`
+        containing the diagonal is returned. If `a` is a `matrix`, a 1-D
         array containing the diagonal is returned in order to maintain
-        backward compatibility.  If the dimension of `a` is greater than
-        two, then an array of diagonals is returned, "packed" from
-        left-most dimension to right-most (e.g., if `a` is 3-D, then the
-        diagonals are "packed" along rows).
+        backward compatibility.
+        If ``a.ndim > 2``, then the dimensions specified by `axis1` and `axis2`
+        are removed, and a new axis inserted at the end corresponding to the
+        diagonal.
 
     Raises
     ------
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 1f70d78e0..67a813bf7 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -4819,7 +4819,7 @@ class MaskedArray(ndarray):
             return result.astype(dtype)
         else:
             D = self.diagonal(offset=offset, axis1=axis1, axis2=axis2)
-            return D.astype(dtype).filled(0).sum(axis=None, out=out)
+            return D.astype(dtype).filled(0).sum(axis=-1, out=out)
     trace.__doc__ = ndarray.trace.__doc__
 
     def dot(self, b, out=None, strict=False):
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
index a99b56309..4dd18182c 100644
--- a/numpy/ma/tests/test_core.py
+++ b/numpy/ma/tests/test_core.py
@@ -3550,6 +3550,11 @@ class TestMaskedArrayMathMethods(object):
                                             axis=0))
         assert_equal(np.trace(mX), mX.trace())
 
+        # gh-5560
+        arr = np.arange(2*4*4).reshape(2,4,4)
+        m_arr = np.ma.masked_array(arr, False)
+        assert_equal(arr.trace(axis1=1, axis2=2), m_arr.trace(axis1=1, axis2=2))
+
     def test_dot(self):
         # Tests dot on MaskedArrays.
         (x, X, XX, m, mx, mX, mXX, m2x, m2X, m2XX) = self.d