Merge pull request #60 from josephwillard/issue_50

Add ability to index graph depth.

Author: brandonwillard

symbolic_pymc/tensorflow/printing.py

@@ -14,44 +14,95 @@
 from symbolic_pymc.tensorflow.meta import TFlowMetaOp
 
 
+class DepthExceededException(Exception):
+    pass
+
+
 class TFlowPrinter(object):
     """A printer that indents and keeps track of already printed subgraphs."""
 
-    def __init__(self, formatter, buffer):
+    def __init__(self, formatter, buffer, depth_lower_idx=0, depth_upper_idx=sys.maxsize):
+        # The buffer to which results are printed
         self.buffer = buffer
+        # A function used to pre-process printed results
         self.formatter = formatter
-        self.indentation = ""
+
+        self.depth_count = 0
+        self.depth_lower_idx, self.depth_upper_idx = depth_lower_idx, depth_upper_idx
+
+        # This is the current indentation string
+        if self.depth_lower_idx > 0:
+            self.indentation = "... "
+        else:
+            self.indentation = ""
+
+        # The set of graphs that have already been printed
         self.printed_subgraphs = set()
 
     @contextmanager
     def indented(self, indent):
         pre_indentation = self.indentation
-        if isinstance(indent, int):
-            self.indentation += " " * indent
-        else:
+
+        self.depth_count += 1
+
+        if self.depth_lower_idx < self.depth_count <= self.depth_upper_idx:
             self.indentation += indent
+
         try:
             yield
+        except DepthExceededException:
+            pass
         finally:
             self.indentation = pre_indentation
+            self.depth_count -= 1
 
     def format(self, obj):
         return self.indentation + self.formatter(obj)
 
-    def print(self, obj):
-        self.buffer.write(self.format(obj))
-        self.buffer.flush()
-
-    def println(self, obj):
-        self.buffer.write(self.format(obj) + "\n")
-        self.buffer.flush()
+    def print(self, obj, suffix=""):
+        if self.depth_lower_idx <= self.depth_count < self.depth_upper_idx:
+            self.buffer.write(self.format(obj) + suffix)
+            self.buffer.flush()
+        elif self.depth_count == self.depth_upper_idx:
+            # Only print the cut-off indicator at the first occurrence
+            self.buffer.write(self.format(f"...{suffix}"))
+            self.buffer.flush()
 
+            # Prevent the caller from traversing at this level or higher
+            raise DepthExceededException()
 
-def tf_dprint(obj, printer=None):
-    """Print a textual representation of a TF graph.
+    def println(self, obj):
+        self.print(obj, suffix="\n")
+
+    def subgraph_add(self, obj):
+        if self.depth_lower_idx <= self.depth_count < self.depth_upper_idx:
+            # Only track printed subgraphs when they're actually printed
+            self.printed_subgraphs.add(obj)
+
+    def __repr__(self):  # pragma: no cover
+        return (
+            "TFlowPrinter\n"
+            f"\tdepth_lower_idx={self.depth_lower_idx},\tdepth_upper_idx={self.depth_upper_idx}\n"
+            f"\tindentation='{self.indentation}',\tdepth_count={self.depth_count}"
+        )
+
+
+def tf_dprint(obj, depth_lower=0, depth_upper=10, printer=None):
+    """Print a textual representation of a TF graph. The output roughly follows the format of `theano.printing.debugprint`.
+
+    Parameters
+    ----------
+    obj : Tensorflow object
+        Tensorflow graph object to be represented.
+    depth_lower : int
+        Used to index specific portions of the graph.
+    depth_upper : int
+        Used to index specific portions of the graph.
+    printer : optional
+        Backend used to display the output.
 
-    The output roughly follows the format of `theano.printing.debugprint`.
     """
+
     if isinstance(obj, tf.Tensor):
         try:
             obj.op
@@ -63,7 +114,7 @@ def tf_dprint(obj, printer=None):
             )
 
     if printer is None:
-        printer = TFlowPrinter(str, sys.stdout)
+        printer = TFlowPrinter(str, sys.stdout, depth_lower, depth_upper)
 
     _tf_dprint(obj, printer)
 
@@ -75,28 +126,22 @@ def _tf_dprint(obj, printer):
 
 @_tf_dprint.register(tf.Tensor)
 @_tf_dprint.register(TFlowMetaTensor)
-def _(obj, printer):
+def _tf_dprint_TFlowMetaTensor(obj, printer):
 
     try:
         shape_str = str(obj.shape.as_list())
     except (ValueError, AttributeError):
         shape_str = "Unknown"
 
     prefix = f'Tensor({getattr(obj.op, "type", obj.op)}):{obj.value_index},\tdtype={getattr(obj.dtype, "name", obj.dtype)},\tshape={shape_str},\t"{obj.name}"'
+
     _tf_dprint(prefix, printer)
 
     if isvar(obj.op):
         return
     elif isvar(obj.op.inputs):
         with printer.indented("|  "):
             _tf_dprint(f"{obj.op.inputs}", printer)
-    elif len(obj.op.inputs) > 0:
-        with printer.indented("|  "):
-            if obj not in printer.printed_subgraphs:
-                printer.printed_subgraphs.add(obj)
-                _tf_dprint(obj.op, printer)
-            else:
-                _tf_dprint("...", printer)
     elif obj.op.type == "Const":
         with printer.indented("|  "):
             if isinstance(obj, tf.Tensor):
@@ -110,10 +155,17 @@ def _(obj, printer):
             _tf_dprint(
                 np.array2string(numpy_val, threshold=20, prefix=printer.indentation), printer
             )
+    elif len(obj.op.inputs) > 0:
+        with printer.indented("|  "):
+            if obj in printer.printed_subgraphs:
+                _tf_dprint("...", printer)
+            else:
+                printer.subgraph_add(obj)
+                _tf_dprint(obj.op, printer)
 
 
 @_tf_dprint.register(tf.Operation)
 @_tf_dprint.register(TFlowMetaOp)
-def _(obj, printer):
+def _tf_dprint_TFlowMetaOp(obj, printer):
     for op_input in obj.inputs:
         _tf_dprint(op_input, printer)

tests/tensorflow/test_printing.py

@@ -144,3 +144,70 @@ def test_numpy():
     )
 
     assert std_out.getvalue() == expected_out.lstrip()
+
+
+@run_in_graph_mode
+def test_depth_indexing():
+    """Make sure graph indexing functions as expected."""
+
+    A = tf.compat.v1.placeholder("float", name="A", shape=tf.TensorShape([None, None]))
+    x = tf.compat.v1.placeholder("float", name="x", shape=tf.TensorShape([None, 1]))
+    y = tf.multiply(1.0, x, name="y")
+
+    z = tf.matmul(A, tf.add(y, y, name="x_p_y"), name="A_dot")
+
+    std_out = io.StringIO()
+    with redirect_stdout(std_out):
+        tf_dprint(z, depth_upper=3)
+
+    expected_out = textwrap.dedent(
+        """
+    Tensor(MatMul):0,\tdtype=float32,\tshape=[None, 1],\t"A_dot:0"
+    |  Tensor(Placeholder):0,\tdtype=float32,\tshape=[None, None],\t"A:0"
+    |  Tensor(Add):0,\tdtype=float32,\tshape=[None, 1],\t"x_p_y:0"
+    |  |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    |  |  |  ...
+    |  |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    |  |  |  ...
+    """
+    )
+
+    assert std_out.getvalue() == expected_out.lstrip()
+
+    std_out = io.StringIO()
+    with redirect_stdout(std_out):
+        tf_dprint(z, depth_lower=1)
+
+    expected_out = textwrap.dedent(
+        """
+    ... Tensor(Placeholder):0,\tdtype=float32,\tshape=[None, None],\t"A:0"
+    ... Tensor(Add):0,\tdtype=float32,\tshape=[None, 1],\t"x_p_y:0"
+    ... |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    ... |  |  Tensor(Const):0,\tdtype=float32,\tshape=[],\t"y/x:0"
+    ... |  |  |  1.
+    ... |  |  Tensor(Placeholder):0,\tdtype=float32,\tshape=[None, 1],\t"x:0"
+    ... |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    ... |  |  ...
+    """
+    )
+
+    assert std_out.getvalue() == expected_out.lstrip()
+
+    std_out = io.StringIO()
+    with redirect_stdout(std_out):
+        tf_dprint(z, depth_lower=1, depth_upper=4)
+
+    expected_out = textwrap.dedent(
+        """
+    ... Tensor(Placeholder):0,\tdtype=float32,\tshape=[None, None],\t"A:0"
+    ... Tensor(Add):0,\tdtype=float32,\tshape=[None, 1],\t"x_p_y:0"
+    ... |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    ... |  |  Tensor(Const):0,\tdtype=float32,\tshape=[],\t"y/x:0"
+    ... |  |  |  ...
+    ... |  |  Tensor(Placeholder):0,\tdtype=float32,\tshape=[None, 1],\t"x:0"
+    ... |  Tensor(Mul):0,\tdtype=float32,\tshape=[None, 1],\t"y:0"
+    ... |  |  ...
+    """
+    )
+
+    assert std_out.getvalue() == expected_out.lstrip()