Add optional parameter update_mapping to FunctionGraph

theano/compile/function_module.py

@@ -159,10 +159,22 @@ def std_fgraph(input_specs, output_specs, accept_inplace=False):
 
     """
     orig_inputs = [spec.variable for spec in input_specs]
-    updates = [spec.update for spec in input_specs if spec.update]
+
+    # Extract the updates and the mapping between update outputs and
+    # the updated inputs.
+    updates = []
+    update_mapping = {}
+    out_idx = len(output_specs)
+    for inp_idx in range(len(input_specs)):
+        if input_specs[inp_idx].update:
+            updates.append(input_specs[inp_idx].update)
+            update_mapping[out_idx] = inp_idx
+            out_idx += 1
+
     orig_outputs = [spec.variable for spec in output_specs] + updates
 
-    fgraph = gof.fg.FunctionGraph(orig_inputs, orig_outputs)
+    fgraph = gof.fg.FunctionGraph(orig_inputs, orig_outputs,
+                                  update_mapping=update_mapping)
 
     for node in fgraph.apply_nodes:
         if getattr(node.op, 'destroy_map', None):

theano/gof/fg.py

@@ -109,7 +109,19 @@ class FunctionGraph(utils.object2):
 
     """
 
-    def __init__(self, inputs, outputs, features=None, clone=True):
+    def __init__(self, inputs, outputs, features=None, clone=True,
+                 update_mapping=None):
+        """
+        Create an FunctionGraph which operates on the subgraph bound by the
+        inputs and outputs sets.
+
+        :param inputs: inputs nodes of the graph, usually declared by the user
+        :param outputs: outputs nodes of the graph.
+        :param clone: If true, we will clone the graph. This is
+        useful to remove the constant cache problem.
+        :param update_mapping: dictionnary describing the mapping between
+        the inputs with updates and the outputs corresponding to their updates
+        """
 
         if clone:
             inputs, outputs = graph.clone(inputs, outputs)
@@ -157,6 +169,7 @@ class FunctionGraph(utils.object2):
         self.node_locks = {}
         self.variable_locks = {}
         self.profile = None
+        self.update_mapping = update_mapping
 
     # Setup a Variable #
     def __setup_r__(self, r):