Repository: Theano/Theano

Issue:
Is there a bug when computing gradient of scan with until?
Hi, 

I want to compute the gradient from the scan's output w.r.t the "sequences" input.
I used a "scan_module.until" to early-stop the iteration.
However, when I ran the following lines, the results seem weird:

```
import theano
import theano.tensor as tensor
import theano.scan_module as ts
import numpy

print theano.__version__

x1 = tensor.vector(dtype='float32')
r,_= theano.scan(lambda i, x: (x[i]*x[i], ts.until(tensor.eq(x[i], 5))),
                sequences=tensor.arange(x1.shape[0]),
                non_sequences=[x1])

g = theano.grad(r.sum(), x1)
f = theano.function([x1], [r, g])

x0 = numpy.arange(20).astype('float32')

print x0
print f(x0)
```

The results are:
```
[  0.   1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.  14.  15.  16.  17.  18.  19.]
[array([  0.,   1.,   4.,   9.,  16.,  25.], dtype=float32),
 array([  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0., 0.,   0.,   0.,  28.,  30.,  32.,  34.,  36.,  38.], dtype=float32)]
```

In my view, since the iterations only ran 5 steps before it stops, the gradient should only be computed for the first 5 elements. However,  it did the last 5 elements of the sequence.

Thank you very much!


Discussion:
Yes, that looks like a bug in scan.
I'm not sure when we are going to be able to investigate and fix it, though.
Note, the first step would be to run with mode=DEBUG_MODE. this would allow us to know if the problem is an optimization or the original graph.
Just an update, we found that the grad of scan just don't check if there is an until or not. So mostly, it isn't implemented and don't return an error about that! It show that this is code never used:)

Anyway, we think it should be easy to implement it and someone is working on implementing it and testing it.

Resolving patch:
diff --git a/theano/scan_module/scan_op.py b/theano/scan_module/scan_op.py
--- a/theano/scan_module/scan_op.py
+++ b/theano/scan_module/scan_op.py
@@ -1981,6 +1981,8 @@ def L_op(self, inputs, outs, dC_douts):
                 self.mintaps[self.n_mit_mot]
         else:
             grad_steps = inputs[0]
+        if self.as_while:
+            n_steps = outs[0].shape[0]
 
         # Restrict the number of grad steps according to
         # self.truncate_gradient
@@ -2198,7 +2200,11 @@ def compute_all_gradients(known_grads):
                 dC_dinps_t[dx + self.n_seqs] += dC_dXtm1
         # Construct scan op
         # Seqs
-        outer_inp_seqs = [x[::-1] for x in inputs[1:1 + self.n_seqs]]
+        if self.as_while:
+            # equivalent to x[:n_steps][::-1]
+            outer_inp_seqs = [x[n_steps - 1::-1] for x in inputs[1:1 + self.n_seqs]]
+        else:
+            outer_inp_seqs = [x[::-1] for x in inputs[1:1 + self.n_seqs]]
         for idx in xrange(self.n_mit_mot + self.n_mit_sot):
             mintap = np.min(self.tap_array[idx])
             if idx < self.n_mit_mot:
@@ -2216,7 +2222,11 @@ def compute_all_gradients(known_grads):
             x[:-1][::-1] for x in self.outer_sitsot_outs(outs)]
         for x in self.outer_nitsot_outs(dC_douts):
             if not isinstance(x.type, DisconnectedType):
-                outer_inp_seqs.append(x[::-1])
+                if self.as_while:
+                    # equivalent to x[:n_steps][::-1]
+                    outer_inp_seqs.append(x[n_steps - 1::-1])
+                else:
+                    outer_inp_seqs.append(x[::-1])
 
         if hasattr(inputs[0].tag, 'test_value'):
             # Here we tests that the new scan input sequence all have
@@ -2224,20 +2234,25 @@ def compute_all_gradients(known_grads):
             # fct add and we want to keep it for all Scan op.  This is
             # used in T_Scan.test_grad_multiple_outs_taps to test
             # that.
+            if self.as_while:
+                n = n_steps.tag.test_value
+            else:
+                n = inputs[0].tag.test_value
             for taps, x in zip(self.mitsot_taps(),
                                self.outer_mitsot_outs(outs)):
                 mintap = np.min(taps)
                 if hasattr(x[::-1][:mintap], 'test_value'):
-                    assert (x[::-1][:mintap].tag.test_value.shape[0] ==
-                            inputs[0].tag.test_value)
+                    assert (x[::-1][:mintap].tag.test_value.shape[0] == n)
             for x in self.outer_sitsot_outs(outs):
                 if hasattr(x[::-1][:-1].tag, 'test_value'):
-                    assert (x[::-1][:-1].tag.test_value.shape[0] ==
-                            inputs[0].tag.test_value)
+                    assert (x[::-1][:-1].tag.test_value.shape[0] == n)
             for x in self.outer_nitsot_outs(outs):
                 if hasattr(x[::-1].tag, 'test_value'):
-                    assert (x[::-1].tag.test_value.shape[0] ==
-                            inputs[0].tag.test_value)
+                    if self.as_while:
+                        assert (x[n_steps - 1::-1].tag.test_value.shape[0] ==
+                                n)
+                    else:
+                        assert (x[::-1].tag.test_value.shape[0] == n)
         outer_inp_seqs += [x[::-1][:np.min(taps)]
                            for taps, x in zip(self.mitsot_taps(),
                                               self.outer_mitsot_outs(outs))]
@@ -2536,7 +2551,8 @@ def compute_all_gradients(known_grads):
                         outer_inp_seqs +
                         outer_inp_mitmot +
                         outer_inp_sitsot +
-                        [inputs[0] for _ in xrange(n_nit_sot)] +
+                        [n_steps if self.as_while else inputs[0]
+                         for _ in xrange(n_nit_sot)] +
                         self.outer_shared(inputs) +
                         self.outer_non_seqs(inputs))
 
@@ -2564,7 +2580,19 @@ def compute_all_gradients(known_grads):
             zip(outputs[offset:offset + self.n_seqs],
                 type_outs[offset:offset + self.n_seqs])):
             if t == 'connected':
-                gradients.append(x[::-1])
+                # If the forward scan is in as_while mode, we need to pad
+                # the gradients, so that they match the size of the input
+                # sequences.
+                if self.as_while:
+                    n_zeros = inputs[0] - n_steps
+                    shp = (n_zeros,)
+                    if x.ndim > 1:
+                        shp = shp + x.shape[1:]
+                    z = tensor.zeros(shp, dtype=x.dtype)
+                    x = tensor.concatenate([x[::-1], z], axis=0)
+                    gradients.append(x)
+                else:
+                    gradients.append(x[::-1])
             elif t == 'disconnected':
                 gradients.append(DisconnectedType()())
             elif t == 'through_shared':
@@ -2580,7 +2608,19 @@ def compute_all_gradients(known_grads):
         end = self.n_mit_mot + self.n_mit_sot + self.n_sit_sot
         for p, (x, t) in enumerate(zip(outputs[:end], type_outs[:end])):
             if t == 'connected':
-                gradients.append(x[::-1])
+                # If the forward scan is in as_while mode, we need to pad
+                # the gradients, so that they match the size of the input
+                # sequences.
+                if self.as_while:
+                    n_zeros = inputs[0] - grad_steps
+                    shp = (n_zeros,)
+                    if x.ndim > 1:
+                        shp = shp + x.shape[1:]
+                    z = tensor.zeros(shp, dtype=x.dtype)
+                    x = tensor.concatenate([x[::-1], z], axis=0)
+                    gradients.append(x)
+                else:
+                    gradients.append(x[::-1])
             elif t == 'disconnected':
                 gradients.append(DisconnectedType()())
             elif t == 'through_shared':
@@ -2611,7 +2651,7 @@ def compute_all_gradients(known_grads):
 
         start = len(gradients)
         gradients += [DisconnectedType()()
-                      for x in xrange(self.n_nit_sot)]
+                      for _ in xrange(self.n_nit_sot)]
         begin = end
 
         end = begin + n_sitsot_outs


Test patch:
diff --git a/theano/scan_module/tests/test_scan.py b/theano/scan_module/tests/test_scan.py
--- a/theano/scan_module/tests/test_scan.py
+++ b/theano/scan_module/tests/test_scan.py
@@ -5529,3 +5529,57 @@ def count_up():
 
         _, updates = theano.scan(count_up, n_steps=20)
         func = theano.function(inputs=[inc], outputs=[], updates=updates)
+
+
+class TestGradUntil(unittest.TestCase):
+
+    def setUp(self):
+        self.x = tensor.vector(name='x')
+        self.threshold = tensor.scalar(name='threshold', dtype='int64')
+        self.seq = np.arange(15, dtype=theano.config.floatX)
+        self.numpy_output = self.seq[:7]**2
+        z = np.zeros(8, dtype=theano.config.floatX)
+        self.numpy_gradient = 2 * np.concatenate([self.seq[:7], z], axis=0)
+
+    def test_grad_until(self):
+        r, _ = theano.scan(lambda x, u: (x * x,
+                                         theano.scan_module.until(x > u)),
+                           sequences=self.x,
+                           non_sequences=[self.threshold])
+        g = theano.grad(r.sum(), self.x)
+        f = theano.function([self.x, self.threshold], [r, g])
+        theano_output, theano_gradient = f(self.seq, 5)
+
+        utt.assert_allclose(theano_output, self.numpy_output)
+        utt.assert_allclose(theano_gradient, self.numpy_gradient)
+
+    def test_grad_until_and_truncate(self):
+        n = 3
+        r, _ = theano.scan(lambda x, u: (x * x,
+                                         theano.scan_module.until(x > u)),
+                           sequences=self.x,
+                           non_sequences=[self.threshold],
+                           truncate_gradient=n)
+        g = theano.grad(r.sum(), self.x)
+        f = theano.function([self.x, self.threshold], [r, g])
+        theano_output, theano_gradient = f(self.seq, 5)
+
+        self.numpy_gradient[:7 - n] = 0
+        utt.assert_allclose(theano_output, self.numpy_output)
+        utt.assert_allclose(theano_gradient, self.numpy_gradient)
+
+    def test_grad_until_and_truncate_sequence_taps(self):
+        n = 3
+        r, _ = theano.scan(lambda x, y, u: (x * y,
+                                            theano.scan_module.until(y > u)),
+                           sequences=dict(input=self.x, taps=[-2, 0]),
+                           non_sequences=[self.threshold],
+                           truncate_gradient=n)
+        g = theano.grad(r.sum(), self.x)
+        f = theano.function([self.x, self.threshold], [r, g])
+        theano_output, theano_gradient = f(self.seq, 6)
+
+        # Gradient computed by hand:
+        numpy_grad = np.array([0, 0, 0, 5, 6, 10, 4, 5, 0, 0, 0, 0, 0, 0, 0])
+        numpy_grad = numpy_grad.astype(theano.config.floatX)
+        utt.assert_allclose(theano_gradient, numpy_grad)