Add dtype and NumPy Const printing to TensorFlow debug print

Author: brandonwillard

symbolic_pymc/tensorflow/printing.py

@@ -1,9 +1,14 @@
 import sys
 
+import numpy as np
+import tensorflow as tf
+
 from functools import singledispatch
 from contextlib import contextmanager
 
-import tensorflow as tf
+from unification import isvar
+
+# from tensorflow.python.framework import tensor_util
 
 from symbolic_pymc.tensorflow.meta import TFlowMetaTensor
 from symbolic_pymc.tensorflow.meta import TFlowMetaOp
@@ -77,22 +82,38 @@ def _(obj, printer):
     except (ValueError, AttributeError):
         shape_str = "Unknown"
 
-    prefix = f'Tensor({obj.op.type}):{obj.value_index},\tshape={shape_str}\t"{obj.name}"'
+    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 len(obj.op.inputs) > 0:
+
+    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):
+                numpy_val = obj.eval(session=tf.compat.v1.Session(graph=obj.graph))
+            elif isvar(obj.op.node_def):
+                _tf_dprint(f"{obj.op.node_def}", printer)
+                return
+            else:
+                numpy_val = obj.op.node_def.attr["value"]
+
+            _tf_dprint(
+                np.array2string(numpy_val, threshold=20, prefix=printer.indentation), printer
+            )
 
 
 @_tf_dprint.register(tf.Operation)
 @_tf_dprint.register(TFlowMetaOp)
 def _(obj, printer):
-    prefix = f'Op({obj.type})\t"{obj.name}"'
-    _tf_dprint(prefix, printer)
-    with printer.indented("|  "):
-        for op_input in obj.inputs:
-            _tf_dprint(op_input, printer)
+    for op_input in obj.inputs:
+        _tf_dprint(op_input, printer)

tests/tensorflow/test_printing.py

@@ -7,8 +7,9 @@
 
 from contextlib import redirect_stdout
 
-from tensorflow.python.eager.context import graph_mode
+from unification import var, Var
 
+from symbolic_pymc.tensorflow.meta import mt
 from symbolic_pymc.tensorflow.printing import tf_dprint
 
 from tests.tensorflow import run_in_graph_mode
@@ -23,11 +24,7 @@ def test_eager_mode():
     X_tf = tf.convert_to_tensor(X)
 
     with pytest.raises(ValueError):
-        tf_dprint(X_tf)
-
-    with graph_mode():
-        X_tf = tf.convert_to_tensor(X)
-        tf_dprint(X_tf)
+        _ = tf_dprint(X_tf)
 
 
 @run_in_graph_mode
@@ -47,17 +44,39 @@ def test_ascii_printing():
         tf_dprint(z)
 
     expected_out = textwrap.dedent('''
-    Tensor(MatMul):0,\tshape=[None, 1]\t"A_dot:0"
-    |  Op(MatMul)\t"A_dot"
-    |  |  Tensor(Placeholder):0,\tshape=[None, None]\t"A:0"
-    |  |  Tensor(Add):0,\tshape=[None, 1]\t"x_p_y:0"
-    |  |  |  Op(Add)\t"x_p_y"
-    |  |  |  |  Tensor(Mul):0,	shape=[None, 1]	"y:0"
-    |  |  |  |  |  Op(Mul)\t"y"
-    |  |  |  |  |  |  Tensor(Const):0,\tshape=[]\t"y/x:0"
-    |  |  |  |  |  |  Tensor(Placeholder):0,\tshape=[None, 1]\t"x:0"
-    |  |  |  |  Tensor(Mul):0,\tshape=[None, 1]\t"y:0"
-    |  |  |  |  |  ...
+    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(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 tf.Graph().as_default(), redirect_stdout(std_out):
+        Var._id = 0
+        tt_lv_inputs_mt = mt.Tensor(mt.Operation(var(), var(), var()), 0, var())
+        tt_const_lv_nodedef_mt = mt.Tensor(mt.Operation(mt.Const.op_def, var(), ()), 0, var())
+        tt_lv_op_mt = mt.Tensor(var(), 0, var())
+        test_mt = mt(1) + tt_lv_inputs_mt + tt_const_lv_nodedef_mt + tt_lv_op_mt
+        tf_dprint(test_mt)
+
+    expected_out = textwrap.dedent('''
+    Tensor(AddV2):0,\tdtype=int32,\tshape=~_11,\t"add:0"
+    |  Tensor(AddV2):0,\tdtype=int32,\tshape=~_12,\t"add:0"
+    |  |  Tensor(AddV2):0,\tdtype=int32,\tshape=~_13,\t"add:0"
+    |  |  |  Tensor(Const):0,\tdtype=int32,\tshape=[],\t"Const:0"
+    |  |  |  |  1
+    |  |  |  Tensor(~_15):0,\tdtype=~_3,\tshape=~_14,\t"~_17"
+    |  |  |  |  ~_2
+    |  |  Tensor(Const):0,\tdtype=~_5,\tshape=~_18,\t"~_20"
+    |  |  |  ~_4
+    |  Tensor(~_6):0,\tdtype=~_7,\tshape=~_21,\t"~_22"
     ''')
 
     assert std_out.getvalue() == expected_out.lstrip()
@@ -73,6 +92,49 @@ def test_unknown_shape():
     with redirect_stdout(std_out):
         tf_dprint(A)
 
-    expected_out = 'Tensor(Placeholder):0,\tshape=Unknown\t"A:0"\n'
+    expected_out = 'Tensor(Placeholder):0,\tdtype=float64,\tshape=Unknown,\t"A:0"\n'
+
+    assert std_out.getvalue() == expected_out.lstrip()
+
+
+@run_in_graph_mode
+def test_numpy():
+    """Make sure we can ascii/text print constant tensors with large Numpy arrays."""
+
+    with tf.Graph().as_default():
+        A = tf.convert_to_tensor(np.arange(100))
+
+    std_out = io.StringIO()
+    with redirect_stdout(std_out):
+        tf_dprint(A)
+
+    expected_out = textwrap.dedent('''
+    Tensor(Const):0,\tdtype=int64,\tshape=[100],\t"Const:0"
+    |  [ 0  1  2 ... 97 98 99]
+    ''')
+
+    assert std_out.getvalue() == expected_out.lstrip()
+
+    N = 100
+    np.random.seed(12345)
+    X = np.vstack([np.random.randn(N), np.ones(N)]).T
+
+    with tf.Graph().as_default():
+        X_tf = tf.convert_to_tensor(X)
+
+    std_out = io.StringIO()
+    with redirect_stdout(std_out):
+        tf_dprint(X_tf)
+
+    expected_out = textwrap.dedent('''
+    Tensor(Const):0,\tdtype=float64,\tshape=[100, 2],\t"Const:0"
+    |  [[-0.20470766  1.        ]
+        [ 0.47894334  1.        ]
+        [-0.51943872  1.        ]
+        ...
+        [-0.74853155  1.        ]
+        [ 0.58496974  1.        ]
+        [ 0.15267657  1.        ]]
+    ''')
 
     assert std_out.getvalue() == expected_out.lstrip()