pydotprint now take an optional parameter mode. If a ProfileMode, print time…

theano/printing.py

@@ -9,7 +9,7 @@ from theano import config
 from gof import Op, Apply
 from theano.gof.python25 import any
 from theano.compile import Function, debugmode
-
+from theano.compile.profilemode import ProfileMode
 
 def debugprint(obj, depth=-1, file=None):
     """Print a computation graph to file
@@ -352,13 +352,15 @@ pprint.assign(lambda pstate, r: hasattr(pstate, 'target') and pstate.target is n
 pp = pprint
 
 
-def pydotprint(fct, outfile=os.path.join(config.compiledir,'theano.pydotprint.png'), compact=True):
+def pydotprint(fct, outfile=os.path.join(config.compiledir,'theano.pydotprint.png'), compact=True, mode=None):
     """
     print to a file in png format the graph of op of a compile theano fct.
 
     :param fct: the theano fct returned by theano.function.
     :param outfile: the output file where to put the graph.
     :param compact: if True, will remove intermediate var that don't have name.
+    :param mode: if a ProfileMode, add to each Apply label (s in apply,% in apply in total op time, % in fct time)
+                         Otherwise ignore it
 
     In the graph, box are an Apply Node(the execution of an op) and ellipse are variable.
     If variable have name they are used as the text(if multiple var have the same name, they will be merged in the graph).
@@ -370,6 +372,8 @@ def pydotprint(fct, outfile=os.path.join(config.compiledir,'theano.pydotprint.pn
     blue ellipse are output of the graph
     grey ellipse are var generated by the graph that are not output and are not used.
     """
+    if not isinstance(mode,ProfileMode):
+        mode=None
     try:
         import pydot as pd
     except:
@@ -402,7 +406,18 @@ def pydotprint(fct, outfile=os.path.join(config.compiledir,'theano.pydotprint.pn
         return varstr
     topo = fct.maker.env.toposort()
     def apply_name(node):
-        return str(node.op).replace(':','_')+'    '+str(topo.index(node))
+        prof_str=''
+        if mode:
+            time = mode.apply_time.get((topo.index(node),node.op),-1)
+            #second, % total time in profiler, %fct time in profiler
+            if mode.local_time[0]==0:
+                pt=0
+            else: pt=time*100/mode.local_time[0]
+            if mode.fct_call[fct]==0:
+                pf=0
+            else: pf = time*100/mode.fct_call_time[fct]
+            prof_str='   (%.3fs,%.3f%%,%.3f%%)'%(time,pt,pf)
+        return str(node.op).replace(':','_')+'    '+str(topo.index(node))+prof_str
 
     # Update the inputs that have an update function
     input_update={}