Add to EquilibriumDB and OPT tracks_on_change_inputs

theano/gof/opt.py

@@ -2065,6 +2065,7 @@ class EquilibriumOptimizer(NavigatorOptimizer):
                  optimizers,
                  failure_callback=None,
                  ignore_newtrees=True,
+                 tracks_on_change_inputs=False,
                  max_use_ratio=None,
                  final_optimizers=None,
                  cleanup_optimizers=None):
@@ -2077,6 +2078,7 @@ class EquilibriumOptimizer(NavigatorOptimizer):
         self.global_optimizers = []
         self.final_optimizers = []
         self.cleanup_optimizers = []
+        self.tracks_on_change_inputs = tracks_on_change_inputs
 
         for opt in optimizers:
             if isinstance(opt, LocalOptimizer):
@@ -2223,8 +2225,14 @@ class EquilibriumOptimizer(NavigatorOptimizer):
                         q.remove(node)
                     except ValueError:
                         pass
-
-            u = self.attach_updater(fgraph, importer, pruner, name=self.name)
+            chin = None
+            if self.tracks_on_change_inputs:
+                def chin(node, i, r, new_r, reason):
+                    if node is not current_node and not isinstance(node, str):
+                        q.append(node)
+            u = self.attach_updater(fgraph, importer, pruner,
+                                    chin=chin,
+                                    name=self.name)
             try:
                 while q:
                     node = q.pop()

theano/gof/optdb.py

@@ -244,16 +244,26 @@ class EquilibriumDB(DB):
         optimization application. This could result in less fgraph iterations,
         but this doesn't mean it will be faster globally.
 
+    tracks_on_change_inputs
+        If True, we will re-apply local opt on nodes whose inputs
+        changed during local optimization application. This could
+        result in less fgraph iterations, but this doesn't mean it
+        will be faster globally.
+
     Notes
     -----
     We can put LocalOptimizer and Optimizer as EquilibriumOptimizer
     suppor both.
 
+    It is probably not a good idea to have ignore_newtrees=False and
+    tracks_on_change_inputs=True
+
     """
 
-    def __init__(self, ignore_newtrees=True):
+    def __init__(self, ignore_newtrees=True, tracks_on_change_inputs=False):
         super(EquilibriumDB, self).__init__()
         self.ignore_newtrees = ignore_newtrees
+        self.tracks_on_change_inputs = tracks_on_change_inputs
         self.__final__ = {}
         self.__cleanup__ = {}
 
@@ -281,6 +291,7 @@ class EquilibriumDB(DB):
             opts,
             max_use_ratio=config.optdb.max_use_ratio,
             ignore_newtrees=self.ignore_newtrees,
+            tracks_on_change_inputs=self.tracks_on_change_inputs,
             failure_callback=opt.NavigatorOptimizer.warn_inplace,
             final_optimizers=final_opts,
             cleanup_optimizers=cleanup_opts)