Initial test for infer_shape

theano/tensor/tests/test_basic.py

@@ -7578,6 +7578,8 @@ def test_alloc_diag():
         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,
@@ -7596,6 +7598,22 @@ def test_alloc_diag():
             assert np.all(rediag == test_val)
 
 
+            # Test infer_shape
+            f_shape = theano.function([x], adiag_op(x).shape)
+            theano.printing.debugprint(f_shape.maker.fgraph.outputs[0]) 
+            output_shape = f_shape(test_val)
+            assert diag_arr.shape == output_shape
+            rediag_shape = np.diagonal(
+                np.ones(output_shape),
+                offset=offset,
+                axis1=axis1,
+                axis2=axis2
+            )
+            assert rediag_shape == test_val.shape
+
+
+
+
 class test_numpy_assumptions(unittest.TestCase):
     # Verify that some assumptions Theano makes on Numpy's behavior still hold.
     def test_ndarray_copy(self):