local_merge_optimizer now register correctly the MergeFeature

theano/gof/opt.py

@@ -1165,7 +1165,7 @@ class FromFunctionLocalOptimizer(LocalOptimizer):
             id(self)), file=stream)
 
 
-def local_optimizer(tracks, inplace=False):
+def local_optimizer(tracks, inplace=False, requirements=()):
     def decorator(f):
         """
         WRITEME
@@ -1177,12 +1177,13 @@ def local_optimizer(tracks, inplace=False):
             for t in tracks:
                 if not (isinstance(t, op.Op) or issubclass(t, op.PureOp)):
                     raise ValueError("Tracks are op classes or instances", f.__module__, f.__name__)
-        requirements = ()
+        req = requirements
         if inplace:
             dh_handler = dh.DestroyHandler
-            requirements = (lambda fgraph:
+            req = tuple(requirements) + (
+                lambda fgraph:
                 fgraph.attach_feature(dh_handler()),)
-        rval = FromFunctionLocalOptimizer(f, tracks, requirements)
+        rval = FromFunctionLocalOptimizer(f, tracks, req)
         rval.__name__ = f.__name__
         return rval
     return decorator

theano/tensor/opt.py

@@ -514,8 +514,13 @@ def register_specialize_device(lopt, *tags, **kwargs):
 # 2) after an local optimization being applied, if the
 #    current node is still in the graph, it will continue to the next
 #    local optimizer. So this won't trigger more iteration.
+def add_merge_feature(fgraph):
+    if not hasattr(fgraph, 'merge_feature'):
+        fgraph.attach_feature(theano.gof.opt.MergeFeature())
+
+
 @register_canonicalize('fast_compile', 'merge')
-@gof.local_optimizer(None)
+@gof.local_optimizer(None, requirements=[add_merge_feature])
 def local_merge_optimizer(node):
     if node.fgraph.merge_feature.scheduled:
         ret = merge_optimizer(node.fgraph)