Fix TF tensor name construction and add NodeDef dtype information

NodeDefs can carry dtype information in their "T" attr.  Since we weren't using
that attr, we weren't able to determine dtypes--in some cases--during
the construction of outputs for meta Operations.

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())