Merge pull request #65 from brandonwillard/fix-bad-attr-metatizing

Stop metatizing all node_def.attr values

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
@@ -88,25 +86,26 @@ def make_opdef_sig(cls, opdef):
                 params[name] = new_param
 
         else:
-            params = []
-            for i_name, i_type in input_args:
+            for i_name, i_type in input_args.items():
                 p = Parameter(i_name, Parameter.POSITIONAL_OR_KEYWORD, annotation=i_type)
                 params[i_name] = p
 
             # These are the ambiguities we're attempting to overcome
             # with the `tf.raw_ops` functions above.
-            for a_name, a_type in attrs:
+            for a_name, a_type in attrs.items():
+
                 if a_name == "T":
                     # This is a type value that will most likely be inferred
                     # from/by the inputs.
                     # TODO: We could check for an `allowed_values` attribute.
                     continue
+
                 p = Parameter(
                     a_name,
                     Parameter.POSITIONAL_OR_KEYWORD,
                     # TODO: We could use the `default_value`
                     # attribute.
-                    default=None,
+                    default=Parameter.empty,
                     annotation=a_type,
                 )
                 params[a_name] = p
@@ -183,7 +182,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 +367,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 +404,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

@@ -16,6 +16,7 @@
                                            TFlowMetaOpDef,
                                            TFlowMetaNodeDef,
                                            TFlowOpName,
+                                           MetaOpDefLibrary,
                                            mt)
 
 from tests.tensorflow import run_in_graph_mode
@@ -321,6 +322,47 @@ def test_inputs_remapping():
     assert z_mt.inputs[1].obj == z.op.inputs[2]
 
 
+@pytest.mark.usefixtures("run_with_tensorflow")
+def test_opdef_sig():
+    """Make sure we can construct an `inspect.Signature` object for a protobuf OpDef when its corresponding function isn't present in `tf.raw_ops`."""
+    from tensorflow.core.framework import op_def_pb2
+
+    custom_opdef_tf = op_def_pb2.OpDef()
+    custom_opdef_tf.name = "MyOpDef"
+
+    arg1_tf = op_def_pb2.OpDef.ArgDef()
+    arg1_tf.name = "arg1"
+    arg1_tf.type_attr = "T"
+
+    arg2_tf = op_def_pb2.OpDef.ArgDef()
+    arg2_tf.name = "arg2"
+    arg2_tf.type_attr = "T"
+
+    custom_opdef_tf.input_arg.extend([arg1_tf, arg2_tf])
+
+    attr1_tf = op_def_pb2.OpDef.AttrDef()
+    attr1_tf.name = "T"
+    attr1_tf.type = "type"
+
+    attr2_tf = op_def_pb2.OpDef.AttrDef()
+    attr2_tf.name = "axis"
+    attr2_tf.type = "int"
+    attr2_tf.default_value.i = 1
+
+    custom_opdef_tf.attr.extend([attr1_tf, attr2_tf])
+
+    opdef_sig = MetaOpDefLibrary.make_opdef_sig(custom_opdef_tf)
+
+    import inspect
+    # These are standard inputs
+    assert opdef_sig.parameters['arg1'].default == inspect._empty
+    assert opdef_sig.parameters['arg2'].default == inspect._empty
+    # These are attributes that are sometimes required by the OpDef
+    assert opdef_sig.parameters['axis'].default == inspect._empty
+    # The obligatory tensor name parameter
+    assert opdef_sig.parameters['name'].default is None
+
+
 @pytest.mark.usefixtures("run_with_tensorflow")
 @run_in_graph_mode
 def test_nodedef():
@@ -337,3 +379,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())