added switch op

theano/scalar/basic.py

@@ -441,6 +441,18 @@ class InRange(LogicalComparison):
 inopenrange = InRange(True, True)
 inclosedrange = InRange(False, False)
 
+class Switch(ScalarOp):
+    nin = 3
+    def impl(self, cond, ift, iff):
+        return ift if cond else iff
+    def c_code(self, node, name, (cond, ift, iff), (z, ), sub):
+        return "%(z)s = %(cond)s ? %(ift)s : %(iff)s;" % locals()
+    def grad(self, (cond, ift, iff), (gz, )):
+        return None, switch(cond, gz, 0), switch(cond, 0, gz)
+    def output_types(self, (cond_t, ift_t, iff_t)):
+        return upcast_out(ift_t, iff_t)
+switch = Switch()
+
 ####################
 # BIT-WISE OPERATORS
 ####################

theano/tensor/basic.py

@@ -910,6 +910,15 @@ def neq(a, b):
     """a != b"""
 
 
+##########################
+# Condition
+##########################
+
+@_scal_elemwise
+def switch(cond, ift, iff):
+    """if cond then ift else iff"""
+
+
 ##########################
 # Bit-wise
 ##########################