Add tests for float16 magma optimizations

theano/gpuarray/opt.py

@@ -2183,7 +2183,8 @@ def local_gpu_magma_qr(op, context_name, inputs, outputs):
         return
     op = GpuMagmaQR(complete=True)
     if inputs[0].dtype == 'float16':
-        return op(inputs[0].astype('float32')).astype('float16')
+        outputs = op(inputs[0].astype('float32'))
+        return [o.astype('float16') for o in outputs]
     return op
 
 

theano/gpuarray/tests/test_linalg.py

@@ -7,12 +7,14 @@ from numpy.linalg.linalg import LinAlgError
 
 import theano
 from theano import config
-from theano.gpuarray.linalg import (GpuCholesky, GpuMagmaMatrixInverse,
+from theano.gpuarray.linalg import (GpuCholesky, GpuMagmaCholesky,
+                                    GpuMagmaEigh, GpuMagmaMatrixInverse,
+                                    GpuMagmaQR, GpuMagmaSVD,
                                     cusolver_available, gpu_matrix_inverse,
-                                    gpu_solve, gpu_svd, GpuMagmaCholesky,
-                                    GpuMagmaQR, GpuMagmaEigh)
-from theano.tensor.nlinalg import matrix_inverse, qr, eigh
-from theano.tensor.slinalg import cholesky
+                                    gpu_solve, gpu_svd)
+from theano.tensor.nlinalg import (SVD, Eigh, MatrixInverse, QRFull,
+                                   QRIncomplete, eigh, matrix_inverse, qr)
+from theano.tensor.slinalg import Cholesky, cholesky
 from theano.tests import unittest_tools as utt
 
 from .. import gpuarray_shared_constructor
@@ -216,6 +218,20 @@ class TestMagma(unittest.TestCase):
         if not config.magma.enabled:
             self.skipTest('Magma is not enabled, skipping test')
 
+    def test_magma_opt_float16(self):
+        ops_to_gpu = [(MatrixInverse(), GpuMagmaMatrixInverse),
+                      (SVD(), GpuMagmaSVD),
+                      (QRFull(mode='reduced'), GpuMagmaQR),
+                      (QRIncomplete(mode='r'), GpuMagmaQR),
+                      # TODO: add support for float16 to Eigh numpy
+                      # (Eigh(), GpuMagmaEigh),
+                      (Cholesky(), GpuMagmaCholesky)]
+        for op, gpu_op in ops_to_gpu:
+            A = theano.tensor.matrix("A", dtype="float16")
+            fn = theano.function([A], op(A), mode=mode_with_gpu.excluding('cusolver'))
+            assert any([isinstance(node.op, gpu_op)
+                        for node in fn.maker.fgraph.toposort()])
+
     def test_gpu_matrix_inverse(self):
         A = theano.tensor.fmatrix("A")