Merge pull request #5421 from tfjgeorge/cusolver

theano/gpuarray/tests/test_linalg.py

@@ -7,6 +7,8 @@ import theano
 from theano.tests import unittest_tools as utt
 from .config import mode_with_gpu
 
+from numpy.linalg.linalg import LinAlgError
+
 # Skip tests if cuda_ndarray is not available.
 from nose.plugins.skip import SkipTest
 from theano.gpuarray.linalg import (cusolver_available, gpu_solve)
@@ -16,7 +18,7 @@ if not cusolver_available:
 
 
 class TestCusolver(unittest.TestCase):
-    def run_gpu_solve(self, A_val, x_val):
+    def run_gpu_solve(self, A_val, x_val, A_struct=None):
         b_val = numpy.dot(A_val, x_val)
         b_val_trans = numpy.dot(A_val.T, x_val)
 
@@ -24,14 +26,19 @@ class TestCusolver(unittest.TestCase):
         b = theano.tensor.matrix("b", dtype="float32")
         b_trans = theano.tensor.matrix("b", dtype="float32")
 
-        solver = gpu_solve(A, b)
-        solver_trans = gpu_solve(A, b_trans, trans='T')
+        if A_struct is None:
+            solver = gpu_solve(A, b)
+            solver_trans = gpu_solve(A, b_trans, trans='T')
+        else:
+            solver = gpu_solve(A, b, A_struct)
+            solver_trans = gpu_solve(A, b_trans, A_struct, trans='T')
+
         fn = theano.function([A, b, b_trans], [solver, solver_trans], mode=mode_with_gpu)
         res = fn(A_val, b_val, b_val_trans)
         x_res = numpy.array(res[0])
         x_res_trans = numpy.array(res[1])
-        utt.assert_allclose(x_res, x_val)
-        utt.assert_allclose(x_res_trans, x_val)
+        utt.assert_allclose(x_val, x_res)
+        utt.assert_allclose(x_val, x_res_trans)
 
     def test_diag_solve(self):
         numpy.random.seed(1)
@@ -41,13 +48,24 @@ class TestCusolver(unittest.TestCase):
                                      1)).astype("float32")
         self.run_gpu_solve(A_val, x_val)
 
+    def test_bshape_solve(self):
+        """
+        Test when shape of b (k, m) is such as m > k
+        """
+        numpy.random.seed(1)
+        A_val = numpy.asarray([[2, 0, 0], [0, 1, 0], [0, 0, 1]],
+                              dtype="float32")
+        x_val = numpy.random.uniform(-0.4, 0.4, (A_val.shape[1],
+                                     A_val.shape[1] + 1)).astype("float32")
+        self.run_gpu_solve(A_val, x_val)
+
     def test_sym_solve(self):
         numpy.random.seed(1)
         A_val = numpy.random.uniform(-0.4, 0.4, (5, 5)).astype("float32")
-        A_sym = (A_val + A_val.T) / 2.0
+        A_sym = numpy.dot(A_val, A_val.T)
         x_val = numpy.random.uniform(-0.4, 0.4, (A_val.shape[1],
                                      1)).astype("float32")
-        self.run_gpu_solve(A_sym, x_val)
+        self.run_gpu_solve(A_sym, x_val, 'symmetric')
 
     def test_orth_solve(self):
         numpy.random.seed(1)
@@ -63,3 +81,34 @@ class TestCusolver(unittest.TestCase):
         x_val = numpy.random.uniform(-0.4, 0.4,
                                      (A_val.shape[1], 4)).astype("float32")
         self.run_gpu_solve(A_val, x_val)
+
+    def test_linalgerrsym_solve(self):
+        numpy.random.seed(1)
+        A_val = numpy.random.uniform(-0.4, 0.4, (5, 5)).astype("float32")
+        x_val = numpy.random.uniform(-0.4, 0.4,
+                                     (A_val.shape[1], 4)).astype("float32")
+        A_val = numpy.dot(A_val.T, A_val)
+        # make A singular
+        A_val[:, 2] = A_val[:, 1] + A_val[:, 3]
+
+        A = theano.tensor.matrix("A", dtype="float32")
+        b = theano.tensor.matrix("b", dtype="float32")
+        solver = gpu_solve(A, b, 'symmetric')
+
+        fn = theano.function([A, b], [solver], mode=mode_with_gpu)
+        self.assertRaises(LinAlgError, fn, A_val, x_val)
+
+    def test_linalgerr_solve(self):
+        numpy.random.seed(1)
+        A_val = numpy.random.uniform(-0.4, 0.4, (5, 5)).astype("float32")
+        x_val = numpy.random.uniform(-0.4, 0.4,
+                                     (A_val.shape[1], 4)).astype("float32")
+        # make A singular
+        A_val[:, 2] = 0
+
+        A = theano.tensor.matrix("A", dtype="float32")
+        b = theano.tensor.matrix("b", dtype="float32")
+        solver = gpu_solve(A, b, trans='T')
+
+        fn = theano.function([A, b], [solver], mode=mode_with_gpu)
+        self.assertRaises(LinAlgError, fn, A_val, x_val)