Stop metatizing all node_def.attr values

Fixes #64.

Author: brandonwillard

symbolic_pymc/tensorflow/meta.py

@@ -6,8 +6,6 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 
-from contextlib import suppress
-
 from inspect import Parameter, Signature
 
 from collections import OrderedDict, UserString
@@ -183,7 +181,7 @@ def __hash__(self):
 def _metatize_tf_object(obj):
     try:
         obj = tf.convert_to_tensor(obj)
-    except TypeError:
+    except (TypeError, ValueError):
         raise ValueError("Could not find a TensorFlow MetaSymbol class for {obj}")
 
     if isinstance(obj, tf.Tensor):
@@ -368,9 +366,8 @@ def _protobuf_convert(cls, k, v):
             from google.protobuf.json_format import MessageToDict
             MessageToDict(obj, use_integers_for_enums=True)
         """
-
         if k == "shape":
-            return tensor_shape.as_shape(v.shape)
+            return metatize(tensor_shape.as_shape(v.shape))
         elif k == "dtype":
             return tf.as_dtype(v.type).name
         elif k == "value":
@@ -406,12 +403,6 @@ def __init__(self, op, name, attr, obj=None):
                         continue
 
                 if k != "T" and k in op_param_names:
-                    # XXX: We can't let `metatize` convert NumPy values;
-                    # otherwise, we'll loop endlessly on "Const" Ops.
-                    if k != "value":
-                        with suppress(ValueError):
-                            v = metatize(v)
-
                     self.attr[k] = v
         else:
             self.attr = attr

tests/tensorflow/test_meta.py

@@ -337,3 +337,25 @@ def test_nodedef():
     norm_rv = mt.RandomStandardNormal(mean=0, stddev=1, shape=(1000,), dtype=tf.float32, name=var())
     assert isinstance(norm_rv, TFlowMetaTensor)
     assert norm_rv.dtype == tf.float32
+
+    # We shouldn't be metatizing all parsed `node_def.attr` values; otherwise,
+    # we won't be able to reconstruct corresponding meta Ops using their meta
+    # OpDefs and inputs.
+    x_test = tf.constant([1.8, 2.2], dtype=tf.float32)
+    y_test = tf.dtypes.cast(x_test, dtype=tf.int32, name="y")
+    y_test_mt = mt(y_test)
+
+    # `ytest_mt.inputs` should have two `.attr` values that are Python
+    # primitives (i.e. int and bool); these shouldn't get metatized and break
+    # our ability to reconstruct the object from its rator + rands.
+    assert y_test_mt == y_test_mt.op.op_def(*y_test_mt.inputs)
+
+
+@pytest.mark.usefixtures("run_with_tensorflow")
+@run_in_graph_mode
+def test_metatize():
+    class CustomClass(object):
+        pass
+
+    with pytest.raises(ValueError):
+        mt(CustomClass())