Merge pull request #5358 from Garrett-R/fix_bug_2015

Fix #2015: diag returns ndarray instead of matrix

1 files changed, 11 insertions, 1 deletions

diff --git a/numpy/matrixlib/tests/test_numeric.py b/numpy/matrixlib/tests/test_numeric.py
index 3588de5e6..91dc92d2e 100644
--- a/numpy/matrixlib/tests/test_numeric.py
+++ b/numpy/matrixlib/tests/test_numeric.py
@@ -2,12 +2,22 @@ from __future__ import division, absolute_import, print_function
 
 from numpy.testing import assert_equal, TestCase, run_module_suite
 from numpy.core import ones
-from numpy import matrix
+from numpy import matrix, diagonal, diag
 
 class TestDot(TestCase):
     def test_matscalar(self):
         b1 = matrix(ones((3, 3), dtype=complex))
         assert_equal(b1*1.0, b1)
 
+
+def test_diagonal():
+    b1 = matrix([[1,2],[3,4]])
+    diag_b1 = matrix([[1, 4]])
+
+    assert_equal(b1.diagonal(), diag_b1)
+    assert_equal(diagonal(b1), diag_b1)
+    assert_equal(diag(b1), diag_b1)
+
+
 if __name__ == "__main__":
     run_module_suite()