Fix diag_indices_from and add test.

1 files changed, 7 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py
index afa94396a..d7e61799a 100644
--- a/numpy/lib/tests/test_index_tricks.py
+++ b/numpy/lib/tests/test_index_tricks.py
@@ -1,4 +1,5 @@
 from numpy.testing import *
+import numpy as np
 from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where,
                     ndenumerate, fill_diagonal, diag_indices,
                     diag_indices_from )
@@ -112,6 +113,12 @@ def test_diag_indices():
                   [[0, 0],
                    [0, 1]]]) )
 
+def test_diag_indices_from():
+    x = np.random.random((4, 4))
+    r, c = diag_indices_from(x)
+    assert_array_equal(r, np.arange(4))
+    assert_array_equal(c, np.arange(4))
+
 
 if __name__ == "__main__":
     run_module_suite()