Merge pull request #77 from brandonwillard/fix-tf-name-and-dtype-inference

Fix TF tensor name construction and add NodeDef dtype information

Author: brandonwillard

symbolic_pymc/tensorflow/meta.py

@@ -9,7 +9,6 @@
 from inspect import Parameter, Signature
 
 from collections import OrderedDict
-from collections.abc import Sequence
 
 from functools import partial
 
@@ -220,11 +219,11 @@ def __init__(self, obj=None):
         super().__init__(obj=obj)
         self._apply_func_sig, self._apply_func = op_def_lib.get_op_info(obj)
 
-    def out_meta_types(self, inputs=None):
+    def out_meta_types(self, inputs=None, node_def=None):
         def _convert_outputs(o):
-            if o.type_attr == "T":
-                return (TFlowMetaTensor, var())
-            elif o.type_attr == "dtype":
+            if o.type_attr == "T" and node_def:
+                return (TFlowMetaTensor, node_def.attr.get("T", var()))
+            elif o.type_attr == "dtype" and inputs:
                 return (TFlowMetaTensor, inputs.get("dtype", var()))
             else:
                 return (TFlowMetaTensor, var())
@@ -284,7 +283,6 @@ def __call__(self, *args, **kwargs):
             apply_arguments.get(i.name) for i in self.obj.input_arg if i.name in apply_arguments
         )
 
-        # Get the `OpDef`-instantiating parameters and call them a "node_def".
         node_attr = {a.name: apply_arguments.get(a.name, a) for a in self.obj.attr}
 
         op_name = op_kwargs.get("name", self.obj.name)
@@ -346,6 +344,8 @@ def _protobuf_convert(cls, k, v):
             return metatize(tensor_shape.as_shape(v.shape))
         elif k == "dtype":
             return tf.as_dtype(v.type).name
+        elif k == "T":
+            return tf.as_dtype(v.type).name
         elif k == "value":
             return tensor_util.MakeNdarray(v.tensor)
         else:
@@ -364,22 +364,17 @@ def __init__(self, op, name, attr, obj=None):
         self.name = name if isvar(name) else str(name)
 
         if not isvar(attr):
-            # We want to limit the attributes we'll consider to those that show
-            # up in an OpDef function's signature (e.g. ignore info about
-            # permissible types).
             opdef_sig, _ = op_def_lib.get_op_info(self.op)
-            op_param_names = opdef_sig.parameters.keys()
-
             _attr = dict()
+
             for k, v in attr.items():
                 if isinstance(v, Message):
                     try:
                         v = self._protobuf_convert(k, v)
                     except TypeError:
-                        continue
+                        v = var()
 
-                if k != "T" and k in op_param_names:
-                    _attr[k] = v
+                _attr[k] = v
 
             self.attr = _attr
         else:
@@ -532,11 +527,12 @@ def outputs(self):
         else:
 
             apply_arguments = self.op_def.input_args(*self.inputs, **self.node_def.attr)
-            out_types_mt = self.op_def.out_meta_types(inputs=apply_arguments)
+            out_types_mt = self.op_def.out_meta_types(
+                inputs=apply_arguments, node_def=self.node_def
+            )
 
             mt_outs = tuple(
-                o_type(self, i, var() if o_dtype is None else o_dtype)
-                for i, (o_type, o_dtype) in enumerate(out_types_mt)
+                o_type(self, i, o_dtype) for i, (o_type, o_dtype) in enumerate(out_types_mt)
             )
 
             self._outputs = mt_outs
@@ -574,7 +570,15 @@ def reify(self):
         if isvar(self.node_def):
             return self
 
-        op_attrs, op_attrs_unreified = meta_reify_iter(self.node_def.attr)
+        op_attrs, op_attrs_unreified = meta_reify_iter(
+            # Only use NodeDef attrs that appear in the OpDef's call signature.
+            # Other NodeDef attrs, like dtype and shape, can be computed.
+            {
+                k: v
+                for k, v in self.node_def.attr.items()
+                if k in self.op_def._apply_func_sig.parameters
+            }
+        )
 
         if not (op_inputs_unreified or op_attrs_unreified or MetaSymbol.is_meta(self.name)):
 
@@ -587,6 +591,8 @@ def reify(self):
             tf_out = self.op_def._apply_func(**apply_arguments)
             op_tf = tf_out.op
 
+            # TODO: Update NodeDef attrs?
+
             assert op_tf is not None
             self._obj = op_tf
             return self.obj
@@ -623,14 +629,8 @@ def name(self):
 
         if self.obj is not None and not isinstance(self.obj, Var):
             name = self.obj.name
-        elif (
-            self.op is not None
-            and not isvar(self.op)
-            and not isvar(self.op.name)
-            and not isinstance(self.op.outputs, Sequence)
-        ):
-            out_num = self.op.outputs.index(self)
-            name = f"{self.op.name}:{out_num}"
+        elif isinstance(getattr(self.op, "name", None), str) and not isvar(self.value_index):
+            name = f"{self.op.name}:{self.value_index}"
         else:
             name = var()
 

tests/tensorflow/test_meta.py

@@ -2,7 +2,7 @@
 If you're debugging/running tests manually, it might help to simply
 disable eager execution entirely:
 
-    > tf.compat.v1.disable_eager_execution()
+    tf.compat.v1.disable_eager_execution()
 """
 import pytest
 import numpy as np
@@ -172,6 +172,21 @@ def test_meta_basic():
     assert a_mt.shape.ndims == 2
     assert a_mt.shape == TFlowMetaTensorShape([1, 2])
 
+    # Make sure that names are properly inferred when there are no base objects
+    # to reference
+    with tf.Graph().as_default():
+        one_mt = mt(1.0)
+        log_mt = mt.log(one_mt)
+        assert log_mt.name == 'Log:0'
+        assert log_mt.dtype == tf.float32
+        assert log_mt.op.outputs[0].dtype == tf.float32
+
+        log_mt._name = None
+        one_mt._obj = None
+        log_mt._obj = None
+        assert log_mt.dtype == tf.float32
+        assert log_mt.name == 'Log:0'
+
 
 @pytest.mark.usefixtures("run_with_tensorflow")
 @run_in_graph_mode
@@ -394,7 +409,6 @@ def test_nodedef():
 
     assert 'compute_uv' in node_def_mt.attr
     assert 'full_matrices' in node_def_mt.attr
-    assert 'T' not in node_def_mt.attr
 
     # Some outputs use nodedef information; let's test those.
     norm_rv = mt.RandomStandardNormal(mean=0, stddev=1, shape=(1000,), dtype=tf.float32, name=var())