correct gradient for ARange

theano/tensor/basic.py

@@ -5464,6 +5464,7 @@ class ARange(Op):
 
         inputs = [start, stop, step]
         outputs = [tensor(self.dtype, (False,))]
+
         return Apply(self, inputs, outputs)
 
     @theano.configparser.change_flags(warn_float64='ignore')
@@ -5523,9 +5524,16 @@ class ARange(Op):
         # no gradient through them
         # stop does not affect the output values,
         # just the output shape, so it is disconnected
-        return [start.zeros_like(),
-                DisconnectedType()(),
-                step.zeros_like()]
+
+        if (self.dtype in discrete_dtypes) and (step > 0):
+            return [start.zeros_like(),
+                    DisconnectedType()(),
+                    step.zeros_like()]
+        else:
+            num_steps_taken = (stop-start)/step
+            return [gz.sum(dtype=config.floatX),
+                    stop.zeros_like(dtype=config.floatX),
+                    (gz*arange(num_steps_taken+1)).sum(dtype=config.floatX)]
 
     def R_op(self, inputs, eval_points):
         return [None]

theano/tensor/tests/test_basic.py

@@ -5868,6 +5868,17 @@ class TestARange(unittest.TestCase):
         assert np.all(f(10, 2, 2) == np.arange(10, 2, 2))
         assert np.all(f(0, 0, 1) == np.arange(0, 0, 1))
 
+    def test_grads(self):
+        start, stop, step = fscalars('start', 'stop', 'step')
+        def f(start, stop, step):
+            return ARange(start.type.dtype)(start, stop, step)
+
+        rng = np.random.RandomState(utt.fetch_seed())
+        utt.verify_grad(f, [np.float32(0),
+                            np.float32(5),
+                            np.float32(1)],
+                        rng=rng)
+
     def test_integers(self):
         # Test arange constructor, on integer outputs
         start, stop, step = iscalars('start', 'stop', 'step')