Repository: Theano/Theano

Issue:
Scan GPU test failure: theano.scan_module.tests.test_scan.T_Scan_Gpuarray
In the daily buildbot:

```
======================================================================
FAIL: test_gpu_memory_usage (theano.scan_module.tests.test_scan.T_Scan_Gpuarray)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/tmp/nightly_build_gpuarray/Theano/theano/scan_module/tests/test_scan.py", line 4812, in test_gpu_memory_usage
    assert len(grad_scan_node.outputs) == 2
AssertionError
```

I'm tring to resume the problem here. Mostly this is due to the optimization PushOutScanOutput. There is probably also interaction with Scan.infer_shape!

I did a PR that tried to fix this, but it have other tests failing: https://github.com/Theano/Theano/pull/3797/files. Not sure why those tests fail. This PR change some code in Scan.infer_shape, but not change the behavior. It just make it easier to change the behavior.

https://github.com/Theano/Theano/pull/4013 will add a mechanism that make sure we don't reintroduce old scan in the graph. When merged, we should rebase the fix on top of it before running our tests.

I'll give more information later.



Discussion:

Resolving patch:
diff --git a/theano/gof/graph.py b/theano/gof/graph.py
--- a/theano/gof/graph.py
+++ b/theano/gof/graph.py
@@ -873,7 +873,8 @@ def clone_get_equiv(inputs, outputs, copy_inputs_and_orphans=True, memo=None):
 
 
 def general_toposort(r_out, deps, debug_print=False,
-                     compute_deps_cache=None, deps_cache=None):
+                     compute_deps_cache=None, deps_cache=None,
+                     clients=None):
     """
     WRITEME
 
@@ -886,6 +887,9 @@ def general_toposort(r_out, deps, debug_print=False,
         deps, but that also cache its results in a dict passed as deps_cache.
     deps_cache : dict
         Must be used with compute_deps_cache.
+    clients : dict
+        If a dict is passed it will be filled with a mapping of node
+        -> clients for each node in the subgraph.
 
     Notes
     -----
@@ -924,8 +928,10 @@ def compute_deps_cache(io):
 
     assert isinstance(r_out, (tuple, list, deque))
 
-    reachable, clients = stack_search(deque(r_out), compute_deps_cache,
-                                      'dfs', True)
+    reachable, _clients = stack_search(deque(r_out), compute_deps_cache,
+                                       'dfs', True)
+    if clients is not None:
+        clients.update(_clients)
     sources = deque([r for r in reachable if not deps_cache.get(r, None)])
 
     rset = set()
@@ -935,7 +941,7 @@ def compute_deps_cache(io):
         if node not in rset:
             rlist.append(node)
             rset.add(node)
-            for client in clients.get(node, []):
+            for client in _clients.get(node, []):
                 deps_cache[client] = [a for a in deps_cache[client]
                                       if a is not node]
                 if not deps_cache[client]:
@@ -951,7 +957,7 @@ def compute_deps_cache(io):
     return rlist
 
 
-def io_toposort(inputs, outputs, orderings=None):
+def io_toposort(inputs, outputs, orderings=None, clients=None):
     """
     WRITEME
 
@@ -959,10 +965,13 @@ def io_toposort(inputs, outputs, orderings=None):
     ----------
     inputs : list or tuple of Variable instances
     outputs : list or tuple of Apply instances
-    orderings: dict
+    orderings : dict
         Key: Apply instance. Value: list of Apply instance.
         It is important that the value be a container with a deterministic
         iteration order. No sets allowed!
+    clients : dict
+        If a dict is provided it will be filled with mappings of
+        node->clients for each node in the subgraph that is sorted
 
     """
     # the inputs are used only here in the function that decides what 'predecessors' to explore
@@ -1013,7 +1022,7 @@ def compute_deps(obj):
 
     topo = general_toposort(outputs, deps=compute_deps,
                             compute_deps_cache=compute_deps_cache,
-                            deps_cache=deps_cache)
+                            deps_cache=deps_cache, clients=clients)
     return [o for o in topo if isinstance(o, Apply)]
 
 
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
@@ -56,7 +56,7 @@
 import numpy
 
 import theano
-from theano import tensor
+from theano import tensor, scalar
 from theano.tensor import opt, get_scalar_constant_value, Alloc, AllocEmpty
 from theano import gof
 from theano.compat import OrderedDict
@@ -694,98 +694,16 @@ def process_node(self, fgraph, node):
                          op.inputs, op.outputs, op.info)
 
         new_scan_node = None
-        local_fgraph_topo = theano.gof.graph.io_toposort(op.inputs, op.outputs)
+        clients = {}
+        local_fgraph_topo = theano.gof.graph.io_toposort(args.inner_inputs,
+                                                         args.inner_outputs,
+                                                         clients=clients)
 
         for nd in local_fgraph_topo:
-            if (isinstance(nd.op, theano.tensor.Dot) and
-                    nd.out in args.inner_out_nit_sot):
-                """
-                The following optimization involves pushing out, after the
-                scan, a Dot whose output is nitsot (not feed back to the inner
-                graph) and where one input is one of scan's input with ndim=2
-                and the other is an intermediate variable in the Scan inner
-                graph with ndim=1.
-
-                The Dot product is pushed out of the scan and its inputs are
-                now the original matrix and a new matrix obtained by
-                concatenating the vectors into a matrix.
-
-                """
-                # Ensure that the output of the Dot is used in the outer
-                # graph to avoid apply the optimization needlessly
-                dot_out_nitsot_idx = args.inner_out_nit_sot.index(nd.out)
-                outer_dot_output = args.outer_out_nit_sot[dot_out_nitsot_idx]
-                if len(outer_dot_output.clients) == 0:
-                    continue
-
-                """
-                Validate that one of the inputs is a matrix AND a
-                non-sequence input to scan and that the other input is a
-                vector and either an sequence input to scan or the result
-                of computation in the inner function of scan.
-
-                """
-                valid_inputs = False
-                idx_matrix_input = -1
-                idx_vector_input = -1
-
-                if (nd.inputs[0].ndim == 2 and
-                    (nd.inputs[0] in args.inner_in_non_seqs or
-                     isinstance(nd.inputs[0], tensor.Constant)) and
-                    nd.inputs[1].ndim == 1 and
-                    (nd.inputs[1] in args.inner_in_seqs or
-                     nd.inputs[1] not in args.inner_inputs)):
-
-                    valid_inputs = True
-                    idx_matrix_input = 0
-                    idx_vector_input = 1
-
-                elif (nd.inputs[1].ndim == 2 and
-                      (nd.inputs[1] in args.inner_in_non_seqs or
-                       isinstance(nd.inputs[1], tensor.Constant)) and
-                      nd.inputs[0].ndim == 1 and
-                      (nd.inputs[0] in args.inner_in_seqs or
-                       nd.inputs[0] not in args.inner_inputs)):
-
-                    valid_inputs = True
-                    idx_matrix_input = 1
-                    idx_vector_input = 0
-
-                if valid_inputs:
-                    # The optimization can be applied on the current Dot
-
-                    # Move out of scan the two inputs to the Dot
-                    (outer_vars,
-                     new_scan_node,
-                     new_scan_args) = self.push_out_inner_vars(fgraph,
-                                                               nd.inputs,
-                                                               node, args)
-                    outer_vector_input = outer_vars[idx_vector_input]
-                    outer_matrix_input = outer_vars[idx_matrix_input]
-
-                    # Perform the Dot outside of scan
-                    if idx_matrix_input == 0:
-                        outer_dot_inputs = [outer_vector_input,
-                                            outer_matrix_input.transpose()]
-                        outer_dot_output = theano.tensor.dot(*outer_dot_inputs)
-                    else:  # idx_matrix_input == 1
-                        outer_dot_inputs = [outer_vector_input,
-                                            outer_matrix_input]
-                        outer_dot_output = theano.tensor.dot(*outer_dot_inputs)
-
-                    # Modify the outer graph to add the outer Dot
-                    fgraph.replace_all(
-                        [(new_scan_args.outer_out_nit_sot[dot_out_nitsot_idx],
-                          outer_dot_output)],
-                        reason="scanOp_pushout_output")
-
-                    break
-
-            elif (isinstance(nd.op, theano.tensor.elemwise.Elemwise) and
-                  isinstance(nd.op.nfunc, numpy.ufunc) and
-                  nd.op.nfunc.__name__ == 'add' and
-                  nd.out in args.inner_out_sit_sot and
-                  self.inner_sitsot_only_last_step_used(nd.out, args)):
+            if (isinstance(nd.op, theano.tensor.elemwise.Elemwise) and
+                    isinstance(nd.op.scalar_op, scalar.Add) and
+                    nd.out in args.inner_out_sit_sot and
+                    self.inner_sitsot_only_last_step_used(nd.out, args)):
 
                 # Ensure that one of the input to the add is the output of
                 # the add from a previous iteration of the inner function
@@ -809,7 +727,7 @@ def process_node(self, fgraph, node):
 
                     if (dot_input.owner is not None and
                         isinstance(dot_input.owner.op, theano.tensor.Dot) and
-                        len(dot_input.clients) == 1 and
+                        len(clients[dot_input]) == 1 and
                         dot_input.owner.inputs[0].ndim == 2 and
                         dot_input.owner.inputs[1].ndim == 2 and
                         self.get_outer_ndim(dot_input.owner.inputs[0], args) == 3 and


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
@@ -3120,6 +3120,25 @@ def rnn_fn(_u, _y, _W):
         utt.assert_allclose(vnh0, tnh0, atol=1e-6)
         utt.assert_allclose(vnW, tnW, atol=1e-6)
 
+    def test_pushout_dot(self):
+        W = tensor.matrix('W')
+        h = tensor.matrix('h')
+
+        o, _ = theano.scan(lambda hi, him1, W: (hi, tensor.dot(hi+him1, W)),
+                           outputs_info=[tensor.zeros([h.shape[1]]), None],
+                           sequences=[h],
+                           non_sequences=[W])
+
+        f = theano.function([W, h], o, mode=mode_with_opt)
+
+        scan_nodes = [x for x in f.maker.fgraph.toposort()
+                     if isinstance(x.op,
+                                   theano.scan_module.scan_op.Scan)]
+        assert len(scan_nodes) == 1
+        scan_op = scan_nodes[0].op
+        assert not any(isinstance(n.op, tensor.Dot) for n in
+                       scan_op.fn.maker.fgraph.apply_nodes)
+
     def test_pushout_all(self):
         W1 = tensor.matrix('W1')
         W2 = tensor.matrix('W2')