Merge pull request #83 from brandonwillard/fix-altered-metatized-objects

Fix altered metatized objects

Author: brandonwillard

symbolic_pymc/tensorflow/meta.py

@@ -43,6 +43,14 @@
 
 
 class MetaOpDefLibrary(object):
+    """A singleton-like object that holds correspondences between TF Python API functions and the `OpDef`s they construct.
+
+    It provides a map of `OpDef` names (lower-cased) to the Python API
+    functions in `tensorflow.raw_ops`, as well as `inspect.Signature` objects
+    for said functions so that default values and lists of arguments (keywords
+    included) can be more easily used.
+
+    """
 
     lower_op_name_to_raw = {
         op_name.lower(): op_name
@@ -130,6 +138,12 @@ def make_opdef_sig(cls, opdef, opdef_py_func=None):
 
     @classmethod
     def get_op_info(cls, opdef):
+        """Return the TF Python API function signature for a given `OpDef`.
+
+        Parameter
+        ---------
+           opdef: str or `OpDef` object (meta or base)
+        """
         if isinstance(opdef, str):
             opdef_name = opdef
             opdef = op_def_registry.get(opdef_name)
@@ -189,6 +203,23 @@ def _metatize_tf_eager(obj):
 class TFlowMetaSymbol(MetaSymbol):
     __slots__ = ()
 
+    @classmethod
+    def _metatize(cls, obj):
+
+        res = super()._metatize(obj)
+        res.validate_objs()
+
+        return res
+
+    def validate_objs(self):
+        # If there is no base object associated with the inputs, then we can't
+        # trust a base object associated with this object (e.g. for the case in
+        # which metatize altered a property in an input).
+        for prop in self.rands():
+            if isinstance(prop, MetaSymbol) and prop.obj is None:
+                self.reset()
+                break
+
 
 class OpDefFactoryType(MetaSymbolType):
     __opdefs__ = {}
@@ -272,7 +303,14 @@ def input_args(self, *args, apply_defaults=True, **kwargs):
         return op_args.arguments
 
     def __call__(self, *args, **kwargs):
-        """Create the meta object(s) resulting from an application of this `OpDef`'s implied `Operation`."""
+        """Create the meta object(s) using the TF Python API's operator functions.
+
+        Each meta `OpDef` is associated with a TF Python function
+        (`self._apply_func`) that is used to construct its `Operation`s.
+
+        See `TFlowMetaTensor.operator` and `TFlowMetaTensor.operator`.
+
+        """
 
         apply_arguments = self.input_args(*args, **kwargs)
 
@@ -368,8 +406,7 @@ def __eq__(self, other):
         if not (type(self) == type(other)):
             return False
 
-        if not (self.base == other.base):
-            return False
+        assert self.base == other.base
 
         return self.obj.name == other.obj.name
 
@@ -390,7 +427,7 @@ class TFlowMetaNodeDef(TFlowMetaSymbol):
     def _metatize(cls, obj):
         res = super()._metatize(obj)
 
-        if "node_attrs" in meta._lvar_defaults_enabled:
+        if obj.op != "Const" and "node_attrs" in meta._lvar_defaults_enabled:
             res.attr = var()
 
         if "names" in meta._lvar_defaults_enabled:
@@ -501,8 +538,7 @@ def _metatize(cls, obj):
         ]
         res = cls(*new_args, obj=obj)
 
-        if meta._lvar_defaults_enabled.issuperset(["node_attrs", "names"]):
-            res.reset()
+        res.validate_objs()
 
         return res
 
@@ -688,13 +724,8 @@ class TFlowMetaTensor(TFlowMetaSymbol, MetaVariable):
     @classmethod
     @cachedmethod(lambda cls: tf_metatize_cache)
     def _metatize(cls, obj):
-
-        res = super()._metatize(obj)
-
-        if meta._lvar_defaults_enabled.issuperset(["node_attrs", "names"]):
-            res.reset()
-
-        return res
+        """Cache Tensors specifically."""
+        return super()._metatize(obj)
 
     def __init__(self, op, value_index, dtype, obj=None):
         self.op = metatize(op)
@@ -734,28 +765,44 @@ def name(self):
 
     @property
     def operator(self):
+        """Return the meta OpDef for this tensor.
+
+        Since meta OpDefs are callable (and dispatch to the corresponding TF
+        Python API function), this object called with arguments provided by
+        `TFlowMetaTensor.inputs` recreates the underlying tensor using the TF
+        Python interface.  This approach has advantages over the purely
+        graph-level approach to constructing meta objects, because--when all
+        arguments are reifiable--it allows us to use purely TF means to
+        construct a meta object (i.e. by first constructing the base object and
+        then "metatizing" it).
+
+        Meta objects produced this way result in less unknown information
+        (e.g. dtypes and shapes) and have the same default values as their base
+        object counterparts (e.g. `Operator` names and `NodeDef.attr` values).
+        """
         if self.op is not None and not isvar(self.op):
             return self.op.op_def
 
     @property
     def inputs(self):
-        """Return the tensor's inputs/rands.
+        """Return the inputs necessary to recreate this object using its TF Python API function.
+
+        These inputs differ from `self.op.inputs` primarily in that they
+        contain the `node_def` parameters as keywords (e.g. to Python API
+        functions like `tf.add`).
 
-        NOTE: These inputs differ from `self.op.inputs` in that they contain
-        the `node_def` parameters, as well.
-        In other words, these can be used to recreate this object (per
-        the meta object spec).
+        See `TFlowMetaTensor.operator` for more information.
         """
         # TODO: In keeping with our desire to return logic variables in cases
         # where params aren't given/inferred, we could return something like
         # `cons(var(), var())` here (although that wouldn't be necessarily imply
         # that the result is a proper list/tuple).
-        if self.op is not None and not isvar(self.op):
-            input_args = self.op.op_def.input_args(
-                *self.op.inputs,
-                name=self.op.name if not isvar(self.op.name) else None,
-                **self.op.node_def.attr,
-            )
+        if self.op is not None and not isvar(self.op) and not isvar(self.op.inputs):
+            if not isvar(self.op.node_def) and not isvar(self.op.node_def.attr):
+                attr = self.op.node_def.attr
+            else:
+                attr = {}
+            input_args = self.op.op_def.input_args(*self.op.inputs, name=self.op.name, **attr)
             return tuple(input_args.values())
 
     def reify(self):

tests/tensorflow/test_meta.py

@@ -64,6 +64,9 @@ def test_meta_eager():
 @run_in_graph_mode
 def test_meta_basic():
 
+    assert mt.Add == mt.Add
+    assert mt.Add != mt.Sub
+
     var_mt = TFlowMetaTensor(var(), var(), var())
     # It should generate a logic variable for the name and use from here on.
     var_name = var_mt.name
@@ -248,9 +251,10 @@ def test_meta_lvars():
     assert all(isvar(getattr(tn_mt, s)) for s in tn_mt.__all_props__)
     assert isinstance(tn_mt.reify(), TFlowMetaTensor)
 
-    mo_mt = TFlowMetaOp(mt.Add, [tn_mt, tn_mt], var())
+    mo_mt = TFlowMetaOp(mt.Add, var(), [tn_mt, var('a')])
     assert len(mo_mt.outputs) == 1
     assert isinstance(mo_mt.reify(), TFlowMetaOp)
+    assert mo_mt.outputs[0].inputs == (tn_mt, var('a'), mo_mt.name)
 
 
 @pytest.mark.usefixtures("run_with_tensorflow")
@@ -593,3 +597,12 @@ def test_global_options():
         assert isvar(b_mt.name)
         assert isvar(b_mt.op.node_def.attr)
         assert b_mt.op.inputs[1] is a_mt
+
+        # `NodeDef.attr` for constants should not be turned into lvars
+        assert not isvar(b_mt.op.inputs[0].op.node_def.attr)
+        assert not isvar(b_mt.op.inputs[1].op.node_def.attr)
+
+    # Make sure we clear out the `.obj` so that the names won't mismatch
+    with tf.Graph().as_default(), enable_lvar_defaults('names'):
+        a_mt = mt(1.0)
+        assert isvar(a_mt.name)