PerformLinker.make_all() can accepts storage_map

theano/gof/link.py

@@ -545,7 +545,7 @@ def map_storage(fgraph, order, input_storage, output_storage, storage_map = None
         input_storage = [[None] for input in fgraph.inputs]
     else:
         assert len(fgraph.inputs) == len(input_storage)
-        
+
     if storage_map is None:
         storage_map = {}
 
@@ -759,7 +759,7 @@ class PerformLinker(LocalLinker):
         self.no_recycling = no_recycling
         return self
 
-    def make_all(self, input_storage=None, output_storage=None):
+    def make_all(self, input_storage=None, output_storage=None, storage_map=None):
         """
 
         Parameters
@@ -781,7 +781,7 @@ class PerformLinker(LocalLinker):
         order = self.schedule(fgraph)
         no_recycling = self.no_recycling
 
-        input_storage, output_storage, storage_map = map_storage(fgraph, order, input_storage, output_storage)
+        input_storage, output_storage, storage_map = map_storage(fgraph, order, input_storage, output_storage, storage_map)
 
         compute_map = {}
         for k in storage_map: