Repository: Theano/Theano

Issue:
Optimization failure (SeqOptimizer, PushOutNonSeqScan)
The following code generates an optimization failure.

```
import theano
outputs, _ = theano.scan(lambda x: (x * x, x), non_sequences=[1], n_steps=2)
f = theano.function(inputs=[], outputs=outputs)
```

When run I see 6 occurrences of this warning/error sequence:

```
WARNING: An optimization wanted to replace a Variable in the graph, but the replacement for it doesn't remove it. We disabled the optimization. Your function runs correctly, but it would be appreciated if you submit this problem to the mailing list theano-users so that we can fix it.
scanOp_pushout_nonseqs_ops [(for{cpu,scan_fn}.0, Alloc.0)]
ERROR (theano.gof.opt): SeqOptimizer apply <theano.scan_module.scan_opt.PushOutNonSeqScan object at 0x0000000004C33FD0>
ERROR (theano.gof.opt): Traceback:
ERROR (theano.gof.opt): Traceback (most recent call last):
  File "<snip>\theano\theano\gof\opt.py", line 194, in apply
    sub_prof = optimizer.optimize(fgraph)
  File "<snip>\theano\theano\gof\opt.py", line 79, in optimize
    ret = self.apply(fgraph, *args, **kwargs)
  File "<snip>\theano\theano\scan_module\scan_opt.py", line 228, in apply
    self.process_node(fgraph, node)
  File "<snip>\theano\theano\scan_module\scan_opt.py", line 401, in process_node
    reason='scanOp_pushout_nonseqs_ops')
  File "<snip>\theano\theano\gof\toolbox.py", line 324, in replace_all_validate_remove
    raise ReplacementDidntRemovedError()
ReplacementDidntRemovedError
```

Then a single instance of:

```
ERROR (theano.gof.opt): EquilibriumOptimizer max'ed out by 'all_pushout_opt'. You can safely raise the current threshold of 5.000000 with the theano flag 'optdb.max_use_ratio'.
```

A few observations:
- Setting `n_steps=1` avoids the problem.
- Returning just `x` from the lambda avoids the problem.
- Returning just `x * x` from the lambda avoids the problem.
- Returning `(x + x, x)` from the lambda also causes the problem.
- Returning `(x, x + x)` from the lambda also causes the problem.



Discussion:

Resolving patch:
diff --git a/theano/scan_module/scan_opt.py b/theano/scan_module/scan_opt.py
--- a/theano/scan_module/scan_opt.py
+++ b/theano/scan_module/scan_opt.py
@@ -397,10 +397,24 @@ def add_to_replace(y):
             # because the scan op expects for a tensor3, to which an
             # subtensor is applied that takes only the last element
             if replace_with:
-                fgraph.replace_all_validate_remove(
-                    replace_with.items(),
-                    remove=[node],
-                    reason='scanOp_pushout_nonseqs_ops')
+                if len(node.outputs) == len(replace_with):
+                    # Every output of the node has a replacement, the Scan
+                    # node can be removed from the graph
+                    fgraph.replace_all_validate_remove(
+                        replace_with.items(),
+                        remove=[node],
+                        reason='scanOp_pushout_nonseqs_ops')
+                else:
+                    # The node has some outputs for which no replacement has
+                    # been established. This can occur for outputs that are
+                    # not produced by apply nodes (since the optimizations
+                    # only visits apply nodes) such as constants or inputs
+                    # passed directly as outputs. The replacements can be
+                    # performed but the Scan node can't be removed at this
+                    # point.
+                    fgraph.replace_all_validate(
+                        replace_with.items(),
+                        reason='scanOp_pushout_nonseqs_ops')
 
         else:
             return False


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
@@ -2771,6 +2771,22 @@ def test_pushout_seqs2(self):
         # an exception being raised
         theano.function([x], outputs, updates=updates)
 
+    @theano.configparser.change_flags(on_opt_error='raise')
+    def test_pushout_nonseq(self):
+        # Test case originally reported by Daniel Renshaw. The crashed occured
+        # during the optimization PushOutNonSeqScan when it attempted to
+        # a scan node with two outputs but only providing a replacement for
+        # one of those outputs. This led the optimization to raise an
+        # exception.
+
+        outputs, _ = theano.scan(lambda x: (x * x, x),
+                                 non_sequences=[2], n_steps=2)
+        f = theano.function(inputs=[], outputs=outputs)
+
+        outs = f()
+        expected_outs = [[4, 4], [2, 2]]
+        utt.assert_allclose(outs, expected_outs)
+
     def test_sequence_dict(self):
         # Test that we can specify sequences as a dictionary with
         # only the 'input' key