Fix test and make it work

theano/gof/tests/test_vm.py

@@ -349,8 +349,8 @@ def test_reallocation():
     x = tensor.scalar('x')
     y = tensor.scalar('y')
     z = tensor.tanh(3 * x + y) + tensor.cosh(x + 5 * y)
-
-    m = theano.compile.get_mode(theano.Mode(linker='vm_nogc'))
+    for l in ['vm_nogc', 'vm', 'vm_nogc', 'vm']:
+        m = theano.compile.get_mode(theano.Mode(linker=l))
         m = m.excluding('fusion', 'inplace')
 
         f = theano.function([x, y], z, name="test_reduce_memory",
@@ -367,8 +367,10 @@ def test_reallocation():
                     keys_copy.remove(i)
                     for o in keys_copy:
                         if (storage_map[i][0] and
-                            storage_map[i][0] == storage_map[o][0]):
+                            storage_map[i][0] is storage_map[o][0]):
                             return [True, storage_map[o][0]]
             return [False, None]
 
         assert check_storage(storage_map)[0]
+        assert len(set([id(v) for v in
+                        storage_map.values()])) < len(storage_map)

theano/gof/vm.py

@@ -1006,7 +1006,7 @@ class VM_Linker(link.LocalLinker):
             lazy = not all([(not th.lazy) for th in thunks])
         if not (lazy or (config.profile and config.profile_memory) or self.use_cloop or self.callback):
             for pair in reallocated_info.values():
-                storage_map[pair[1]][0] = storage_map[pair[0]][0]
+                storage_map[pair[1]] = storage_map[pair[0]]
 
         computed, last_user = link.gc_helper(order)
         if self.allow_gc: