improving debugprint for scan ops

theano/compile/debugmode.py

@@ -521,7 +521,8 @@ def debugprint(r, prefix='', depth=-1, done=None, print_type=False,
                file=sys.stdout, print_destroy_map=False,
                print_view_map=False, order=None, ids='CHAR',
                stop_on_name=False, prefix_child=None,
-               scan_ops=None, profile=None):
+               scan_ops=None, profile=None,
+               scan_inner_to_outer_inputs=None):
     """Print the graph leading to `r` to given depth.
 
     :param r: Variable instance
@@ -544,6 +545,8 @@ def debugprint(r, prefix='', depth=-1, done=None, print_type=False,
                          we don't print anything below it.
     :param scan_ops: Scan ops in the graph will be added inside this list
                      for later printing purposes.
+    :param scan_inner_to_outer_inputs: a dictionary mapping a scan ops inner function
+    inputs to the scan op inputs (outer inputs) for printing purposes.
 
     """
     if depth == 0:
@@ -578,6 +581,7 @@ def debugprint(r, prefix='', depth=-1, done=None, print_type=False,
         elif ids == "":
             id_str = ""
         done[obj] = id_str
+
         return id_str
 
     if hasattr(r.owner, 'op'):
@@ -681,9 +685,21 @@ def debugprint(r, prefix='', depth=-1, done=None, print_type=False,
                     debugprint(i, new_prefix, depth=depth - 1, done=done,
                                print_type=print_type, file=file, order=order,
                                ids=ids, stop_on_name=stop_on_name,
-                               prefix_child=new_prefix_child,
-                               scan_ops=scan_ops, profile=profile)
+                               prefix_child=new_prefix_child, scan_ops=scan_ops,
+                               profile=profile,
+                               scan_inner_to_outer_inputs=scan_inner_to_outer_inputs)
+    else:
+        if scan_inner_to_outer_inputs is not None and\
+           r in scan_inner_to_outer_inputs:
 
+            id_str = get_id_str(r)
+            outer_r = scan_inner_to_outer_inputs[r]
+
+            if hasattr(outer_r.owner, 'op'):
+                outer_id_str = get_id_str(outer_r.owner)
+            else:
+                outer_id_str = get_id_str(outer_r)
+            print('%s%s %s%s -> %s' % (prefix, r, id_str, type_str, outer_id_str), file=file)
         else:
             # this is an input variable
             id_str = get_id_str(r)

theano/printing.py

@@ -149,6 +149,7 @@ N.B.:
                              file=_file, order=order, ids=ids,
                              scan_ops=scan_ops, stop_on_name=stop_on_name,
                              profile=p)
+
     if len(scan_ops) > 0:
         print("", file=_file)
         new_prefix = ' >'
@@ -156,17 +157,31 @@ N.B.:
         print("Inner graphs of the scan ops:", file=_file)
 
         for s in scan_ops:
+            # prepare a dict which maps the scan op's inner inputs to its outer inputs.
+            if hasattr(s.owner.op, 'fn'):
+                # If the op was compiled, print the optimized version.
+                inner_inputs = s.owner.op.fn.maker.fgraph.inputs
+            else:
+                inner_inputs = s.owner.op.inputs
+            outer_inputs = s.owner.inputs
+            inner_to_outer_inputs = dict([(inner_inputs[i],outer_inputs[o])
+                                          for i,o in enumerate(
+                                                  s.owner.op.get_outer_iidx_from_inner_iidx_seq())])
+            #import pdb; pdb.set_trace()
+
             print("", file=_file)
             debugmode.debugprint(s, depth=depth, done=done,
                                  print_type=print_type,
                                  file=_file, ids=ids,
-                                 scan_ops=scan_ops, stop_on_name=stop_on_name)
+                                 scan_ops=scan_ops, stop_on_name=stop_on_name,
+                                 scan_inner_to_outer_inputs=inner_to_outer_inputs)
             if hasattr(s.owner.op, 'fn'):
                 # If the op was compiled, print the optimized version.
                 outputs = s.owner.op.fn.maker.fgraph.outputs
             else:
                 outputs = s.owner.op.outputs
             for idx, i in enumerate(outputs):
+
                 if hasattr(i, 'owner') and hasattr(i.owner, 'op'):
                     if isinstance(i.owner.op, theano.scan_module.scan_op.Scan):
                         scan_ops.append(i)
@@ -176,7 +191,8 @@ N.B.:
                                      print_type=print_type, file=_file,
                                      ids=ids, stop_on_name=stop_on_name,
                                      prefix_child=new_prefix_child,
-                                     scan_ops=scan_ops)
+                                     scan_ops=scan_ops,
+                                     scan_inner_to_outer_inputs=inner_to_outer_inputs)
 
     if file is _file:
         return file