Initial draft of test.

theano/tensor/tests/test_basic.py

@@ -7561,6 +7561,34 @@ class test_diag(unittest.TestCase):
         tensor.verify_grad(diag, [x], rng=rng)
 
 
+def test_alloc_diag(self):
+    dims = 4
+    shape = (5,) * dims
+    xv = np.random.randn(*shape).astype(config.floatX)
+    for d in xrange(1, dims + 1):
+        # Create a TensorType of the same dimensions as
+        # as the data we want to test.
+        x = T.TensorType(dtype=config.floatX, broadcastable=(False,) * d)('x')
+        # Make a slice of the test data that has the
+        # dimensions we need by doing xv[0,...,0]
+        # For example, for an array of shape (5,), we
+        # need to do xv[0, 0, 0, 0].
+        test_val = xv[((0,) * (dims - d))]
+        for offset, axis1, axis2 in [(0, 0, 1),]:
+            adiag_op = AllocDiag(offset=offset,
+                                 axis1=axis1,
+                                 axis2=axis2)
+            f = theano.function([x], adiag_op(x))
+            # AllocDiag and extract the diagonal again
+            # to check
+            rediag = np.diagonal(
+                f(xv),
+                offset=offset,
+                axis1=axis1,
+                axis2=axis2
+            )
+            assert (rediag == x).all()
+
 class test_numpy_assumptions(unittest.TestCase):
     # Verify that some assumptions Theano makes on Numpy's behavior still hold.
     def test_ndarray_copy(self):