test the inplace optimization for ifelse

theano/tests/test_ifelse.py

@@ -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 .. even if the two results have different
+        # shapes but the same type (i.e. both vectors, or matrices or
+        # whatnot of same dtype
+        x = tensor.vector('x')
+        y = tensor.vector('y')
+        c = tensor.iscalar('c')
+        f = theano.function([c, x, y], 2*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.inplace for x in f.maker.env.toposort() if
+                         x.op == IfElse)
+        assert numpy.allclose(vx*2, f(1, vx, vy))
+        assert numpy.allclose(vy*2, f(0, vx, vy))
+
     def test_lazy_if_on_generics(self):
         x = theano.generic()
         y = theano.generic()