Fix IfElse gpu optimization that used the old ifelse interface.

theano/sandbox/cuda/opt.py

@@ -11,6 +11,7 @@ import theano
 from theano.scan_module import scan_utils, scan_op, scan_opt
 from theano import scalar as scal
 from theano import tensor, compile, gof
+import theano.ifelse
 
 from theano.compile import optdb
 from theano.gof import (local_optimizer, EquilibriumDB, SequenceDB, ProxyDB,
@@ -366,36 +367,45 @@ def local_gpu_lazy_ifelse(node):
 
     ifelse(host_from_gpu) -> host_from_gpu(ifelse)
     """
-    if hasattr(theano, "lazycond"):
-        gpu_ifelse = theano.lazycond.IfElse(gpu=True)
-
-        if node.op == gpu_from_host:
-            host_input = node.inputs[0]
-            if (host_input.owner
-                    and host_input.owner.op == theano.lazycond.ifelse):
-                c, t, f = host_input.owner.inputs
-                if not isinstance(f.type, CudaNdarrayType):
-                    f = gpu_from_host(f)
-                if not isinstance(t.type, CudaNdarrayType):
-                    t = gpu_from_host(t)
-                if isinstance(c.type, CudaNdarrayType):
-                    c = host_from_gpu(c)
-
-                return [gpu_ifelse(c, t, f)]
-
-        if node.op == theano.lazycond.ifelse:
-            if numpy.any([(i.owner and i.owner.op == host_from_gpu)
-                          for i in node.inputs]):
-                c, t, f = node.inputs
-
-                if not isinstance(f.type, CudaNdarrayType):
-                    f = gpu_from_host(f)
-                if not isinstance(t.type, CudaNdarrayType):
-                    t = gpu_from_host(t)
-                if isinstance(c.type, CudaNdarrayType):
-                    c = host_from_gpu(c)
-
-                return [host_from_gpu(gpu_ifelse(c, t, f))]
+    if isinstance(node.op, theano.ifelse.IfElse) and not node.op.gpu:
+        gpu_ifelse = theano.ifelse.IfElse(node.op.n_outs, gpu=True)
+        if numpy.any([(i.owner and i.owner.op == host_from_gpu)
+                      for i in node.inputs]) or numpy.any(
+                      [c != 'output' and c.op == gpu_from_host for c, idx
+                       in node.outputs[0].clients]):
+
+            c = node.inputs[0]
+            outs = node.inputs[1:]
+            # Should not happen, but just in case
+            if isinstance(c.type, CudaNdarrayType):
+                c = host_from_gpu(c)
+
+            for i in range(len(outs)):
+                if not isinstance(outs[i], CudaNdarrayType):
+                    outs[i] = gpu_from_host(outs[i])
+            return [host_from_gpu(out) for out in
+                    gpu_ifelse.make_node(c, *outs).outputs]
+
+    if node.op == gpu_from_host:
+        host_input = node.inputs[0]
+        if (host_input.owner and
+            isinstance(host_input.owner.op, theano.ifelse.IfElse) and
+            not host_input.owner.op.gpu):
+            gpu_ifelse = theano.ifelse.IfElse(host_input.owner.op.n_outs,
+                                                  gpu=True)
+
+            c = host_input.owner.inputs[0]
+            outs = host_input.owner.inputs[1:]
+            # Should not happen, but just in case
+            if isinstance(c.type, CudaNdarrayType):
+                c = host_from_gpu(c)
+
+            for i in range(len(outs)):
+                if not isinstance(outs[i], CudaNdarrayType):
+                    outs[i] = gpu_from_host(outs[i])
+
+            outs = gpu_ifelse.make_node(c, *outs).outputs
+            return outs
 
     return False