Merge pull request #3152 from carriepl/scan_aswhile_opt

theano/scan_module/scan_opt.py

@@ -612,8 +612,13 @@ class PushOutScanOutput(gof.Optimizer):
         fgraph.attach_feature(gof.toolbox.ReplaceValidate())
 
     def apply(self, fgraph):
+        # Don't perform the optimization on as_while scans. Because these scans
+        # don't run for a predetermined number of steps, handling them is
+        # more complicated and this optimization doesn't support it at the
+        # moment.
         nodelist = [x for x in fgraph.toposort()
-                    if isinstance(x.op, scan_op.Scan)]
+                    if (isinstance(x.op, scan_op.Scan) and
+                        not x.op.as_while)]
         for node in nodelist:
             # Process the node as long as something gets optimized
             while node != None:

theano/scan_module/tests/test_scan.py

@@ -3029,6 +3029,43 @@ class T_Scan(unittest.TestCase):
         sol[:, :] = v_out
         utt.assert_allclose(sol, f(v_h, v_W1, v_W2))
 
+    def test_pushout_while(self):
+        # Ensure that the optimizations for Scan that push computation out of
+        # the Scan don't alter the result for 'as_while' scans.
+
+        W1 = tensor.matrix('W1')
+        W2 = tensor.matrix('W2')
+        step_indices = tensor.vector('step_indices')
+
+        def lambda_fn(step_idx, W1, W2):
+            until_condition = theano.scan_module.until(step_idx > 2)
+            return tensor.dot(W1, W2), until_condition
+
+        # Compile a function with the optimization
+        o, _ = theano.scan(lambda_fn,
+                           sequences=[step_indices, W1],
+                           non_sequences=[W2],
+                           n_steps=5)
+
+        f = theano.function([W1, W2, step_indices], o, mode=mode_with_opt)
+
+        # Compule an theano function without the optimization
+        o, _ = theano.scan(lambda_fn,
+                           sequences=[step_indices, W1],
+                           non_sequences=[W2],
+                           n_steps=5, mode='FAST_COMPILE')
+
+        f_ref = theano.function([W1, W2, step_indices], o, mode='FAST_COMPILE')
+
+        # Compare the results of the two implementations
+        input_values = [numpy.random.random((5, 5)).astype("float32"),
+                        numpy.random.random((5, 5)).astype("float32"),
+                        numpy.arange(5).astype("float32")]
+
+        out = f(*input_values)
+        out_ref = f_ref(*input_values)
+        utt.assert_allclose(out, out_ref)
+
     def test_pushout(self):
         W1 = tensor.matrix('W1')
         W2 = tensor.matrix('W2')