Modify link.map_storage() so that it make use of given storage_map

theano/gof/link.py

@@ -496,10 +496,17 @@ class Container(object):
         return r
 
 
-def map_storage(fgraph, order, input_storage, output_storage):
-    """
-    Ensure there is storage (a length-1 list) for inputs, outputs, and
-    interior nodes.
+
+def map_storage(fgraph, order, input_storage, output_storage, storage_map = None):
+    """Ensure there is storage (a length-1 list) for inputs, outputs, and interior nodes.
+
+    :param fgraph: The current fgraph.  This function uses the inputs and outputs attributes.
+    :param order: an iterable over Apply instances (in program running order)
+    :param input_storage: None or existing input storage (see below)
+    :param output_storage: None or existing output storage (see below)
+
+    :rtype: 3-tuple
+    :returns: (list of storage for inputs, list of storage for outputs, and the `storage_map`)
 
     Parameters
     ----------
@@ -538,8 +545,9 @@ def map_storage(fgraph, order, input_storage, output_storage):
         input_storage = [[None] for input in fgraph.inputs]
     else:
         assert len(fgraph.inputs) == len(input_storage)
-
-    storage_map = {}
+        
+    if storage_map is None:
+        storage_map = {}
 
     # add input storage into storage_map
     for r, storage in zip(fgraph.inputs, input_storage):