Merge pull request #245 from pascanur/test_inplace_ifelse

theano/ifelse.py

@@ -351,10 +351,11 @@ def cond_make_inplace(node):
     return False
 
 
+optdb.register('cond_make_inplace', opt.in2out(cond_make_inplace,
+    ignore_newtrees=True), 95, 'fast_run', 'inplace')
+
 # XXX: Optimizations commented pending further debugging (certain optimizations
 # make computation less lazy than it should be currently).
-# optdb.register('cond_make_inplace', opt.in2out(cond_make_inplace,
-#     ignore_newtrees=True), 95, 'fast_run', 'inplace')
 #
 # ifelse_equilibrium = gof.EquilibriumDB()
 # ifelse_seqopt = gof.SequenceDB()

theano/tests/test_ifelse.py

@@ -21,7 +21,7 @@ class test_ifelse(unittest.TestCase):
     def test_lazy_if(self):
         # Tests that lazy if works .. even if the two results have different
         # shapes but the same type (i.e. both vectors, or matrices or
-        # whatnot of same dtype
+        # whatnot of same dtype)
         x = tensor.vector('x')
         y = tensor.vector('y')
         c = tensor.iscalar('c')
@@ -37,6 +37,26 @@ class test_ifelse(unittest.TestCase):
         assert numpy.allclose(vx, f(1, vx, vy))
         assert numpy.allclose(vy, f(0, vx, vy))
 
+    def test_lazy_if_inplace(self):
+        # Tests that lazy if works inplace
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        f = theano.function([c, x, y], ifelse(c, x, y))
+        rng = numpy.random.RandomState(utt.fetch_seed())
+
+        xlen = rng.randint(200)
+        ylen = rng.randint(200)
+
+        vx = numpy.asarray(rng.uniform(size=(xlen,)), theano.config.floatX)
+        vy = numpy.asarray(rng.uniform(size=(ylen,)), theano.config.floatX)
+        assert numpy.all([x.op.as_view for x in f.maker.env.toposort() if
+                         isinstance(x.op, IfElse)])
+        assert len([x.op for x in f.maker.env.toposort()
+                   if isinstance(x.op, IfElse)]) > 0
+        assert numpy.allclose(vx, f(1, vx, vy))
+        assert numpy.allclose(vy, f(0, vx, vy))
+
     def test_lazy_if_on_generics(self):
         x = theano.generic()
         y = theano.generic()