Now mrg.normal accept a theano variable as size.

theano/sandbox/rng_mrg.py

@@ -10,7 +10,7 @@ import numpy
 
 from theano import Op, Apply, shared, config, Variable
 from theano.tensor import raw_random, TensorType, as_tensor_variable, get_vector_length, cast, opt
-from theano.tensor import zeros_like, sqrt, log, sin, cos, join
+from theano.tensor import zeros_like, sqrt, log, sin, cos, join, prod
 from theano.compile import optdb
 from theano.gof import local_optimizer
 
@@ -612,7 +612,7 @@ class MRG_RandomStreams(object):
 
     def n_streams(self, size):
         # TODO: a smart way of choosing the number of streams
-        if isinstance(size, (tuple, list)):
+        if isinstance(size, (tuple, list)) and all([isinstance(i,int) for i in size]):
             r = 1
             for s in size:
                 r *= s
@@ -639,13 +639,12 @@ class MRG_RandomStreams(object):
         ndim may be a plain integer to supplement the missing
         information.
         
-        :param: size: Can be a list of integer or a Theano variable like the shape of some tensor.
-                      The number of dimensions must be computable at compile time.
+        :param: size: Can be a list of integer or Theano variable(ex: the shape of other Theano Variable)
                       TODO: can size be None?
         """
         if isinstance(size, tuple):
-            assert all([isinstance(i,int) for i in size]), "size must be a tuple of int or a Theano variable"
-        else: assert isinstance(size, Variable), "size must be a tuple of int or a Theano variable"
+            assert all([isinstance(i,int) or isinstance(i,Variable) for i in size]), "size must be a tuple of int or a Theano variable"
+        else: assert isinstance(size, Variable) and size.ndim==1, "size must be a tuple of int or a Theano variable"
 
         if nstreams is None:
             nstreams = self.n_streams(size)
@@ -706,24 +705,33 @@ class MRG_RandomStreams(object):
             raise NotImplementedError("MRG_RandomStreams.multinomial only implemented with n == 1 and pvals.ndim = 2")
 
     def normal(self, size=None, avg=0.0, std=1.0, ndim=None, dtype=config.floatX):
+        """
+        :param: size: Can be a list of integer or Theano variable(ex: the shape of other Theano Variable)
+        """
         # We need an even number of ]0,1[ samples. Then we split them
         # in two halves. First half becomes our U1's for Box-Muller,
         # second half our U2's. See Wikipedia page:
         # http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
 
-        assert isinstance(size, tuple), "size must be a tuple"
-        assert all([isinstance(i,int) for i in size])
-        n_samples = numpy.prod(size)
         evened = False
+        constant = False
+        if isinstance(size, tuple) and all([isinstance(i,int) for i in size]):
+            constant = True          
+            n_samples = numpy.prod(size)
            
-        if n_samples % 2 == 1:
-            n_samples += 1
-            evened = True
-
+            if n_samples % 2 == 1:
+                n_samples += 1
+                evened = True
+        else:
+            n_samples = prod(size)+(prod(size)%2)#if even, don't change, if odd, +1
         flattened = self.uniform(size=(n_samples,), dtype=dtype)
 
-        U1 = flattened[:n_samples/2]
-        U2 = flattened[n_samples/2:]
+        if constant:
+            U1 = flattened[:n_samples/2]
+            U2 = flattened[n_samples/2:]
+        else:
+            U1 = flattened[:prod(flattened.shape)/2]
+            U2 = flattened[prod(flattened.shape)/2:]
 
         #normal_samples = zeros_like(flattened)
         sqrt_ln_U1 = sqrt(-2.0*log(U1))
@@ -740,8 +748,10 @@ class MRG_RandomStreams(object):
         final_samples = None
         if evened:
             final_samples = normal_samples[:-1]
-        else:
+        elif constant:
             final_samples = normal_samples
+        else:
+            final_samples = normal_samples[:prod(size)]
 
         final_samples = avg + std * final_samples
 

theano/sandbox/test_rng_mrg.py

@@ -433,53 +433,55 @@ def test_normal0():
 
     steps = 50
     if mode in ['DEBUG_MODE','FAST_COMPILE']:
-        sample_size = (99,30)
+        sample_size = (25,30)
         rtol=.02
     else:
         sample_size = (999,50)
         rtol=.01
+    sample_size_odd = (sample_size[0],sample_size[1]-1)
+    x = tensor.matrix()
+    for size,const_size,var_input,input in [(sample_size,sample_size,[],[]), (x.shape,sample_size,[x],[numpy.zeros(sample_size)]),
+                                 (sample_size_odd,sample_size_odd,[],[]),#test odd value
+                                 (x.shape,sample_size_odd,[x],[numpy.zeros(sample_size_odd)]),#test odd value
+                                 ]:
+        print ''
+        print 'ON CPU:'
 
-    print ''
-    print 'ON CPU:'
-
-    R = MRG_RandomStreams(234, use_cuda=False)
-    n = R.normal(size=sample_size, avg=-5.0, std=2.0)
-    f = theano.function([], n, mode=mode)
-    theano.printing.debugprint(f)
-    print 'random?[:10]\n', f()[0,0:10]
-    basictest(f, steps, sample_size, target_avg=-5.0, target_std=2.0, prefix='mrg ', allow_01=True, mean_rtol=rtol)
+        R = MRG_RandomStreams(234, use_cuda=False)
+        n = R.normal(size=size, avg=-5.0, std=2.0)
+        f = theano.function(var_input, n, mode=mode)
+        theano.printing.debugprint(f)
+        print 'random?[:10]\n', f(*input)[0,0:10]
+        basictest(f, steps, const_size, target_avg=-5.0, target_std=2.0, prefix='mrg ', allow_01=True, inputs=input, mean_rtol=rtol)
 
-    sys.stdout.flush()
+        sys.stdout.flush()
 
-    # now with odd number of samples
-    sample_size = (sample_size[0],sample_size[1]-1)
+        if mode!='FAST_COMPILE' and cuda_available:
+            print ''
+            print 'ON GPU:'
+            R = MRG_RandomStreams(234, use_cuda=True)
+            n = R.normal(size=size, avg=-5.0, std=2.0, dtype='float32')
+            assert n.dtype == 'float32' #well, it's really that this test w GPU doesn't make sense otw
+            f = theano.function(var_input, theano.Out(
+                theano.sandbox.cuda.basic_ops.gpu_from_host(n),
+                borrow=True), mode=mode_with_gpu)
 
-    if mode!='FAST_COMPILE' and cuda_available:
-        print ''
-        print 'ON GPU:'
-        R = MRG_RandomStreams(234, use_cuda=True)
-        n = R.normal(size=sample_size, avg=-5.0, std=2.0, dtype='float32')
-        assert n.dtype == 'float32' #well, it's really that this test w GPU doesn't make sense otw
-        f = theano.function([], theano.Out(
-            theano.sandbox.cuda.basic_ops.gpu_from_host(n),
-            borrow=True), mode=mode_with_gpu)
+            theano.printing.debugprint(f)
+            sys.stdout.flush()
+            print 'random?[:10]\n', numpy.asarray(f(*input))[0,0:10]
+            print '----'
+            sys.stdout.flush()
+            basictest(f, steps, const_size, target_avg=-5.0, target_std=2.0, prefix='gpu mrg ', allow_01=True, inputs=input, mean_rtol=rtol)
 
-        theano.printing.debugprint(f)
-        sys.stdout.flush()
-        print 'random?[:10]\n', numpy.asarray(f())[0,0:10]
-        print '----'
-        sys.stdout.flush()
-        basictest(f, steps, sample_size, target_avg=-5.0, target_std=2.0, prefix='gpu mrg ', allow_01=True, mean_rtol=rtol)
-        
 
-    print ''
-    print 'ON CPU w NUMPY:'
-    RR = theano.tensor.shared_randomstreams.RandomStreams(234)
+        print ''
+        print 'ON CPU w NUMPY:'
+        RR = theano.tensor.shared_randomstreams.RandomStreams(234)
 
-    nn = RR.normal(size=sample_size, avg=-5.0, std=2.0)
-    ff = theano.function([], nn)
+        nn = RR.normal(size=size, avg=-5.0, std=2.0)
+        ff = theano.function(var_input, nn)
 
-    basictest(ff, steps, sample_size, target_avg=-5.0, target_std=2.0, prefix='numpy ', allow_01=True, mean_rtol=rtol)
+        basictest(ff, steps, const_size, target_avg=-5.0, target_std=2.0, prefix='numpy ', allow_01=True, inputs=input, mean_rtol=rtol)
 
 def basic_multinomialtest(f, steps, sample_size, target_pvals, prefix="", mean_rtol=0.04):