Test _infer_ndim_bcast in more cases, use it instead of _infer_ndim.

theano/tensor/raw_random.py

@@ -235,10 +235,6 @@ class RandomFunction(gof.Op):
         return [None for i in inputs]
 
 
-def _infer_ndim(ndim, shape, *args):
-    ndim, ivec, bcast = _infer_ndim_bcast(ndim, shape, *args)
-    return ndim, ivec
-
 def _infer_ndim_bcast(ndim, shape, *args):
     """
     Infer the number of dimensions from the shape or the other arguments.
@@ -402,10 +398,10 @@ def uniform(random_state, size=None, low=0.0, high=1.0, ndim=None, dtype=theano.
     """
     low = tensor.as_tensor_variable(low)
     high = tensor.as_tensor_variable(high)
-    ndim, size = _infer_ndim(ndim, size, low, high)
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size, low, high)
     dtype = tensor.scal.upcast(dtype, low.dtype, high.dtype)
     op = RandomFunction('uniform',
-            tensor.TensorType(dtype = dtype, broadcastable = (False,)*ndim) )
+            tensor.TensorType(dtype=dtype, broadcastable=bcast) )
     return op(random_state, size, low, high)
 
 def binomial(random_state, size=None, n=1, p=0.5, ndim=None, dtype='int64', prob=None):
@@ -424,7 +420,7 @@ def binomial(random_state, size=None, n=1, p=0.5, ndim=None, dtype='int64', prob
         print >> sys.stderr, "DEPRECATION WARNING: the parameter prob to the binomal fct have been renamed to p to have the same name as numpy."
     n = tensor.as_tensor_variable(n)
     p = tensor.as_tensor_variable(p)
-    ndim, size = _infer_ndim(ndim, size, n, p)
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size, n, p)
     if n.dtype=='int64':
         ### THIS WORKS AROUND A NUMPY BUG on 32bit machine
         ###  Erase when the following works on a 32bit machine:
@@ -449,10 +445,10 @@ def normal(random_state, size=None, avg=0.0, std=1.0, ndim=None, dtype=theano.co
     """
     avg = tensor.as_tensor_variable(avg)
     std = tensor.as_tensor_variable(std)
-    ndim, size = _infer_ndim(ndim, size, avg, std)
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size, avg, std)
     dtype = tensor.scal.upcast(dtype, avg.dtype, std.dtype)
     op = RandomFunction('normal',
-            tensor.TensorType(dtype = dtype, broadcastable = (False,)*ndim) )
+            tensor.TensorType(dtype=dtype, broadcastable=bcast))
     return op(random_state, size, avg, std)
 
 def random_integers_helper(random_state, low, high, size):
@@ -509,9 +505,9 @@ def random_integers(random_state, size=None, low=0, high=1, ndim=None, dtype='in
     """
     low = tensor.as_tensor_variable(low)
     high = tensor.as_tensor_variable(high)
-    ndim, size = _infer_ndim(ndim, size, low, high)
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size, low, high)
     op = RandomFunction(random_integers_helper,
-            tensor.TensorType(dtype = dtype, broadcastable = (False,)*ndim) )
+            tensor.TensorType(dtype=dtype, broadcastable=bcast))
     return op(random_state, size, low, high)
 
 def permutation_helper(random_state, n, shape):
@@ -558,10 +554,10 @@ def permutation(random_state, size=None, n=1, ndim=None, dtype='int64'):
     .. note::
         Note that the output will then be of dimension ndim+1.
     """
-    ndim, size = _infer_ndim(ndim, size)
+    ndim, size, bcast = _infer_ndim_bcast(ndim, size)
     #print "NDIM", ndim, size
     op = RandomFunction(permutation_helper,
-            tensor.TensorType(dtype=dtype, broadcastable=(False,)*(ndim+1)),
+            tensor.TensorType(dtype=dtype, broadcastable=bcast+(False,)),
             ndim_added=1)
     return op(random_state, size, n)
 

theano/tensor/tests/test_raw_random.py

@@ -165,6 +165,51 @@ class T_random_function(unittest.TestCase):
         self.assertTrue(numpy.allclose(o4, o1_4))
         self.assertTrue(numpy.allclose(o4, o2_4_4[0]))
 
+    def test_random_function_noshape_args(self):
+        '''Test if random_function helper works with args but without shape'''
+        rng_R = random_state_type()
+
+        # No shape, default args -> OK
+        post_out, out = uniform(rng_R, size=None, ndim=2)
+        f = compile.function(
+                [compile.In(rng_R,
+                    value=numpy.random.RandomState(utt.fetch_seed()),
+                    update=post_out,
+                    mutable=True)],
+                [out],
+                accept_inplace=True)
+        o, = f()
+
+        # No shape, args that have to be broadcasted -> OK
+        low = tensor.TensorType(dtype='float64',
+                broadcastable=(False, True, True))()
+        high = tensor.TensorType(dtype='float64',
+                broadcastable=(True, True, True, False))()
+        post_out2, out2 = uniform(rng_R, size=None, ndim=2, low=low, high=high)
+        self.assertEqual(out2.ndim, 4)
+        self.assertEqual(out2.broadcastable, (True,False,True,False))
+
+        g = compile.function(
+                [low,
+                 high,
+                 compile.In(rng_R,
+                    value=numpy.random.RandomState(utt.fetch_seed()),
+                    update=post_out2,
+                    mutable=True)],
+                [out2],
+                accept_inplace=True)
+        low_v = [[[3]],[[4]],[[-5]]]
+        high_v = [[[[5, 8]]]]
+        o2, = g(low_v, high_v)
+        self.assertEqual(o2.shape, (1,3,1,2))
+
+    def test_random_function_noshape_noargs(self):
+        '''Test if random_function helper works without args or shape'''
+        rng_R = random_state_type()
+
+        # No shape, no args -> TypeError
+        self.assertRaises(TypeError, permutation, rng_R, size=None, ndim=2)
+
     def test_random_function_ndim_added(self):
         """Test that random_function helper function accepts ndim_added as keyword argument"""
         # If using numpy's uniform distribution, ndim_added should be 0,
@@ -174,10 +219,14 @@ class T_random_function(unittest.TestCase):
         # and a ValueError should be raised.
         def ndim_added_deco(ndim_added):
             def randomfunction(random_state, size=(), low=0.0, high=0.0, ndim=None):
-                ndim, size = raw_random._infer_ndim(ndim, size)
+                ndim, size, bcast = raw_random._infer_ndim_bcast(ndim, size)
+                if ndim_added < 0:
+                    bcast = bcast[:ndim_added]
+                else:
+                    bcast = bcast + ((False,) * ndim_added)
+                assert len(bcast) == ndim + ndim_added
                 op = RandomFunction('uniform',
-                        tensor.TensorType(dtype = 'float64', broadcastable =
-                            (False,)*(ndim+ndim_added)),
+                        tensor.TensorType(dtype='float64', broadcastable=bcast),
                         ndim_added=ndim_added)
                 return op(random_state, size, low, high)
             return randomfunction