Refactored AllocDiag tests for both GPU and CPU.

theano/gpuarray/tests/test_subtensor.py

@@ -5,7 +5,7 @@ import unittest
 import theano
 from theano import tensor
 from theano.compile import DeepCopyOp
-from theano.tensor.tests import test_subtensor
+from theano.tensor.tests import test_subtensor, test_basic
 from theano.tests import unittest_tools as utt
 
 from ..basic_ops import HostFromGpu, GpuFromHost, GpuContiguous
@@ -318,6 +318,15 @@ class test_gpuextractdiag(unittest.TestCase):
                 np_x.diagonal(offset, axis1, axis2))
 
 
+class test_gpu_alloc_diag(test_basic.test_alloc_diag):
+    def __init__(self, name):
+        return test_basic.test_alloc_diag.__init__(
+            self, name,
+            alloc_diag=GpuAllocDiag,
+            mode=mode_with_gpu
+        )
+
+
 class test_gpuallocdiag(unittest.TestCase):
     def test_allocdiag_opt(self):
         x = tensor.vector()

theano/tensor/tests/test_basic.py

@@ -7561,7 +7561,17 @@ class test_diag(unittest.TestCase):
         tensor.verify_grad(diag, [x], rng=rng)
 
 
-def test_alloc_diag():
+class test_alloc_diag(unittest.TestCase):
+    def __init__(self, name, alloc_diag=AllocDiag, mode=None):
+        self.alloc_diag = AllocDiag
+
+        if mode is None:
+            mode = theano.compile.mode.get_default_mode()
+        self.mode = mode
+
+        return super(test_alloc_diag, self).__init__(name)
+
+    def _generator(self):
         dims = 4
         shape = (5,) * dims
         xv = np.random.randn(*shape).astype(config.floatX)
@@ -7575,14 +7585,17 @@ def test_alloc_diag():
             # For example, for an array of shape (5,), we
             # need to do xv[0, 0, 0, 0].
             test_val = xv[((0,) * (dims - d))]
+            yield x, test_val
+
+    def test_alloc_diag_values(self):
+        for x, test_val in self._generator():
             for offset, axis1, axis2 in [(0, 0, 1), (0, 1, 2), (1, 0, 1),
                                          (0, 1, 3), (0, 2, 3), (1, 2, 3),
                                          (-1, 0, 1), (-2, 0, 1), (-1, 1, 2)]:
-
                 # Test AllocDiag values
                 if np.maximum(axis1, axis2) > len(test_val.shape):
                     continue
-            adiag_op = AllocDiag(offset=offset,
+                adiag_op = self.alloc_diag(offset=offset,
                                            axis1=axis1,
                                            axis2=axis2)
                 f = theano.function([x], adiag_op(x))
@@ -7602,7 +7615,7 @@ def test_alloc_diag():
 
                 theano.printing.debugprint(f_shape.maker.fgraph.outputs[0])
                 output_shape = f_shape(test_val)
-            assert not any(isinstance(node.op, AllocDiag)
+                assert not any(isinstance(node.op, self.alloc_diag)
                                for node in f_shape.maker.fgraph.toposort())
                 rediag_shape = np.diagonal(
                     np.ones(output_shape),
@@ -7612,6 +7625,23 @@ def test_alloc_diag():
                 ).shape
                 assert np.all(rediag_shape == test_val.shape)
 
+                diag_x = adiag_op(x)
+                sum_diag_x = tensor.sum(diag_x)
+                grad_x = tensor.grad(sum_diag_x, x)
+                grad_diag_x = tensor.grad(sum_diag_x, diag_x)
+                f_grad_x = theano.function([x], grad_x, mode=self.mode)
+                f_grad_diag_x = theano.function([x], grad_diag_x, mode=self.mode)
+                grad_input = f_grad_x(test_val)
+                grad_diag_input = f_grad_diag_x(test_val)
+                true_grad_input = np.diagonal(
+                    grad_diag_input,
+                    offset=offset,
+                    axis1=axis1,
+                    axis2=axis2
+                )
+
+                assert np.all(true_grad_input == grad_input)
+
 
 class test_numpy_assumptions(unittest.TestCase):
     # Verify that some assumptions Theano makes on Numpy's behavior still hold.