Fix test and make it work

theano/gof/tests/test_vm.py

@@ -349,26 +349,28 @@ 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'))
-    m = m.excluding('fusion', 'inplace')
-
-    f = theano.function([x, y], z, name="test_reduce_memory",
-                        mode=m)
-    output = f(1, 2)
-    assert output
-    storage_map = f.fn.storage_map
-
-    def check_storage(storage_map):
-        from theano.tensor.var import TensorConstant
-        for i in storage_map.keys():
-            if not isinstance(i, TensorConstant):
-                keys_copy = storage_map.keys()[:]
-                keys_copy.remove(i)
-                for o in keys_copy:
-                    if (storage_map[i][0] and
-                            storage_map[i][0] == storage_map[o][0]):
-                        return [True, storage_map[o][0]]
-        return [False, None]
-
-    assert check_storage(storage_map)[0]
+    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",
+                            mode=m)
+        output = f(1, 2)
+        assert output
+        storage_map = f.fn.storage_map
+
+        def check_storage(storage_map):
+            from theano.tensor.var import TensorConstant
+            for i in storage_map.keys():
+                if not isinstance(i, TensorConstant):
+                    keys_copy = storage_map.keys()[:]
+                    keys_copy.remove(i)
+                    for o in keys_copy:
+                        if (storage_map[i][0] and
+                            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: