allow ones/zeros to accept single integer shapes (plus tests)

theano/tensor/basic.py

@@ -2064,6 +2064,8 @@ def zeros(shape, dtype=None):
     """
     Create a Tensor filled with zeros, closer to Numpy's syntax than ``alloc``.
     """
+    if not isinstance(shape, (list, tuple)):
+        shape = [shape]
     if dtype is None:
         dtype = config.floatX
     return alloc(numpy.array(0, dtype=dtype), *shape)
@@ -2073,6 +2075,8 @@ def ones(shape, dtype=None):
     """
     Create a Tensor filled with ones, closer to Numpy's syntax than ``alloc``.
     """
+    if not isinstance(shape, (list, tuple)):
+        shape = [shape]
     if dtype is None:
         dtype = config.floatX
     return alloc(numpy.array(1, dtype=dtype), *shape)

theano/tensor/tests/test_basic.py

@@ -1896,6 +1896,18 @@ class TestAlloc(unittest.TestCase):
                               for node in topo]) == 1
             assert not isinstance(topo[0].op, DeepCopyOp)
 
+    def test_ones(self):
+        shapes = [[], 1, [1], [1, 2], [1, 2, 3]]
+        for shp in shapes:
+            ones = theano.function([], [tensor.ones(shp)])
+            assert numpy.allclose(ones(), numpy.ones(shp))
+
+    def test_zeros(self):
+        shapes = [[], 1, [1], [1, 2], [1, 2, 3]]
+        for shp in shapes:
+            zeros = theano.function([], [tensor.zeros(shp)])
+            assert numpy.allclose(zeros(), numpy.zeros(shp))
+
 
 def test_eye():
     def check(dtype, N, M_=None, k=0):