new_io_toposort takes list of cmp functions

theano/gof/fg.py

@@ -553,8 +553,15 @@ class FunctionGraph(utils.object2):
             # 1-element graphs.
             return list(self.apply_nodes)
         fg = self
+
+        # TODO - This is a hack.
+        #        Should change deeper in to just move around functions
         ords = self.orderings()
-        order = graph.io_toposort(fg.inputs, fg.outputs, ords)
+        ords = [lambda a,b: (1 if b in d.get(a, ())
+                       else -1 if a in d.get(b, ())
+                       else  0)                      for d in ords]
+
+        order = graph.new_io_toposort(fg.inputs, fg.outputs, ords)
         return order
 
     def orderings(self):

theano/gof/graph.py

@@ -1031,6 +1031,14 @@ def dependence(a, b):
         return -1
     return 0
 
-def new_io_toposort(inputs, outputs, cmp=dependence):
+def new_io_toposort(inputs, outputs, cmps=[]):
     """ Same as io_toposort """
+    cmps = [dependence] + cmps # enforce that dependence is the strongest cmp fn
+
+    # An aggregate comparator - looks at each cmp in order
+    def cmp(a,b, fns=cmps):
+        if not fns:  return 0
+        head, tail = fns[0], fns[1:]
+        return head(a, b) or cmp(a, b, tail)
+
     return sorted(list_of_nodes(inputs, outputs), cmp=cmp)