BUG: Ensure lstsq can handle RHS with all sizes.

mhvk · mhvk · commit a3c489f1b252 · 2017-11-06T10:43:37.000-05:00
This fixes a bug in the creation of workspace arrays for a call
to `lapack_lite.zgelsd`, which led to segmentation faults when
a RHS was passed in that had larger size than the size of the
matrix.
diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py
@@ -2021,14 +2021,8 @@ def lstsq(a, b, rcond="warn"):
         work = zeros((lwork,), t)
         results = lapack_routine(m, n, n_rhs, a, m, bstar, ldb, s, rcond,
                                  0, work, -1, rwork, iwork, 0)
-        lwork = int(abs(work[0]))
-        rwork = zeros((lwork,), real_t)
-        a_real = zeros((m, n), real_t)
-        bstar_real = zeros((ldb, n_rhs,), real_t)
-        results = lapack_lite.dgelsd(m, n, n_rhs, a_real, m,
-                                     bstar_real, ldb, s, rcond,
-                                     0, rwork, -1, iwork, 0)
         lrwork = int(rwork[0])
+        lwork = int(work[0].real)
         work = zeros((lwork,), t)
         rwork = zeros((lrwork,), real_t)
         results = lapack_routine(m, n, n_rhs, a, m, bstar, ldb, s, rcond,
diff --git a/numpy/linalg/tests/test_regression.py b/numpy/linalg/tests/test_regression.py
@@ -137,6 +137,18 @@ def test_norm_object_array(self):
         assert_raises(TypeError, linalg.norm, testmatrix, ord=-2)
         assert_raises(ValueError, linalg.norm, testmatrix, ord=3)
 
+    def test_lstsq_complex_larger_rhs(self):
+        # gh-9891
+        size = 20
+        n_rhs = 70
+        G = np.random.randn(size, size) + 1j * np.random.randn(size, size)
+        u = np.random.randn(size, n_rhs) + 1j * np.random.randn(size, n_rhs)
+        b = G.dot(u)
+        # This should work without segmentation fault.
+        u_lstsq, res, rank, sv = linalg.lstsq(G, b, rcond=None)
+        # check results just in case
+        assert_array_almost_equal(u_lstsq, u)
+
 
 if __name__ == '__main__':
     run_module_suite()
