Updates and fix algo

theano/compile/profiling.py

@@ -740,7 +740,7 @@ class ProfileStats(object):
                 Generate all valid node order from node_list
                 and compute its memory peaf
                 """
-                global mem_count, mem_bound, max_mem_count, max_mem_count
+                global mem_count, mem_bound, max_mem_count
 
                 for node in executable_nodes:
                     new_exec_nodes = executable_nodes.copy()
@@ -749,6 +749,8 @@ class ProfileStats(object):
                     mem_created = 0
                     mem_freed = 0
 
+                    max_storage = max_mem_count
+
                     # check if we cut path now
                     if max_mem_count > mem_bound:
                         continue
@@ -776,8 +778,7 @@ class ProfileStats(object):
                     for val in node.inputs:
                         if (dependencies[val] and val.owner and val not in fgraph.outputs):
                             if all(compute_map[v] for v in dependencies[val]):
-                                if not dmap and not vmap:
-                                    mem_freed += var_mem[val]
+                                mem_freed += var_mem[val]
 
                     mem_count -= mem_freed
 
@@ -797,7 +798,7 @@ class ProfileStats(object):
                     
                     # Reset track variables
                     mem_count -= mem_created
-                    max_mem_count -= mem_created
+                    max_mem_count = max_storage
                     mem_count += mem_freed
                     for var in node.outputs:
                         compute_map[var][0] = 0

theano/compile/tests/test_profiling.py

@@ -44,7 +44,7 @@ def test_profiling():
         the_string = buf.getvalue()
         assert "Max if linker=cvm(default): 8208KB (16400KB)" in the_string
         assert "Minimum peak from all valid apply node order is 8192KB" in the_string
-        
+
     finally:
         theano.config.profile = old1
         theano.config.profile_memory = old2