Merge pull request #998 from goodfeli/fix_switch_grad

theano/scalar/basic.py

@@ -1041,15 +1041,13 @@ class Switch(ScalarOp):
         return "%(z)s = %(cond)s ? %(ift)s : %(iff)s;" % locals()
 
     def grad(self, (cond, ift, iff), (gz, )):
-        if ift.type in continuous_types:
-            first_part = switch(cond, gz, 0)
-        else:
-            first_part = None
+        first_part = switch(cond, gz, 0.)
+        second_part = switch(cond, 0., gz)
 
-        if iff.type in continuous_types:
-            second_part = switch(cond, 0, gz)
-        else:
-            second_part = None
+        out = self(cond, ift, iff)
+        if out.type.dtype in discrete_types:
+            first_part = 0.
+            second_part = 0.
 
         # cond does affect the elements of the output so it is connected.
         # For the sake of making the gradient convenient we assume that

theano/scalar/tests/test_basic.py

@@ -214,6 +214,20 @@ def test_grad_gt():
     g = theano.gradient.grad(z, y)
     assert g.eval({ y : 1. }) == 0.
 
+def test_grad_switch():
+
+    # This is a code snippet from the mailing list
+    # It caused an assert to be raised due to the
+    # switch op's grad method not handling integer
+    # inputs correctly
+
+    x = theano.tensor.matrix()
+    c = theano.tensor.matrix()
+
+    s = theano.tensor.switch(c, x, 0)
+    l = s.sum()
+
+    theano.gradient.grad(l, x)
 
 # Testing of Composite is done in tensor/tests/test_opt.py
 # in test_fusion, TestCompositeCodegen