Use the correct MetaOpDef call signatures, function combination

Author: brandonwillard

symbolic_pymc/tensorflow/meta.py

@@ -48,16 +48,14 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @classmethod
-    def make_opdef_sig(cls, opdef):
+    def make_opdef_sig(cls, opdef, opdef_py_func=None):
         """Create a `Signature` object for an `OpDef`.
 
         Annotations are include so that one can partially verify arguments.
         """
         input_args = OrderedDict([(a.name, a.type or a.type_attr) for a in opdef.input_arg])
         attrs = OrderedDict([(a.name, a.type) for a in opdef.attr])
 
-        opdef_py_func = getattr(tf.raw_ops, opdef.name, None)
-
         params = OrderedDict()
         if opdef_py_func:
             # We assume we're dealing with a function from `tf.raw_ops`.
@@ -86,6 +84,9 @@ def make_opdef_sig(cls, opdef):
                 params[name] = new_param
 
         else:
+            # We're crafting the Operation from a low-level via `apply_op`.
+            opdef_py_func = partial(op_def_lib.apply_op, opdef.name)
+
             for i_name, i_type in input_args.items():
                 p = Parameter(i_name, Parameter.POSITIONAL_OR_KEYWORD, annotation=i_type)
                 params[i_name] = p
@@ -117,15 +118,17 @@ def make_opdef_sig(cls, opdef):
         opdef_sig = Signature(
             params.values(), return_annotation=[(o.name, o.type_attr) for o in opdef.output_arg]
         )
-        return opdef_sig
+        return opdef_sig, opdef_py_func
 
     def add_op(self, opdef):
         op_info = self._ops.get(opdef.name, None)
         if op_info is None:
             super().add_op(opdef)
             op_info = self._ops[opdef.name]
-            opdef_sig = self.make_opdef_sig(op_info.op_def)
+            opdef_func = getattr(tf.raw_ops, opdef.name, None)
+            opdef_sig, opdef_func = self.make_opdef_sig(op_info.op_def, opdef_func)
             op_info.opdef_sig = opdef_sig
+            op_info.opdef_func = opdef_func
         return op_info
 
     def get_opinfo(self, opdef):
@@ -239,7 +242,7 @@ class TFlowMetaOpDef(MetaOp, TFlowMetaSymbol):
     def __init__(self, obj=None):
         op_info = op_def_lib.add_op(obj)
         self.apply_func_sig = op_info.opdef_sig
-        self.apply_func = partial(op_def_lib.apply_op, obj.name)
+        self.apply_func = op_info.opdef_func
         super().__init__(obj=obj)
 
     def out_meta_types(self, inputs=None):

tests/tensorflow/test_meta.py

@@ -167,6 +167,7 @@ def test_meta_create():
     with pytest.raises(TypeError):
         TFlowMetaTensor('float64', 'Add', name='q__')
 
+
 @pytest.mark.usefixtures("run_with_tensorflow")
 @run_in_graph_mode
 def test_meta_Op():
@@ -196,7 +197,6 @@ def test_meta_Op():
         assert MetaSymbol.is_meta(test_op.outputs[0])
 
 
-
 @pytest.mark.usefixtures("run_with_tensorflow")
 def test_meta_lvars():
     """Make sure we can use lvars as values."""
@@ -381,7 +381,7 @@ def test_opdef_sig():
 
     custom_opdef_tf.attr.extend([attr1_tf, attr2_tf])
 
-    opdef_sig = MetaOpDefLibrary.make_opdef_sig(custom_opdef_tf)
+    opdef_sig, opdef_func = MetaOpDefLibrary.make_opdef_sig(custom_opdef_tf)
 
     import inspect
     # These are standard inputs
@@ -431,3 +431,13 @@ class CustomClass(object):
 
     with pytest.raises(ValueError):
         mt(CustomClass())
+
+
+@pytest.mark.usefixtures("run_with_tensorflow")
+@run_in_graph_mode
+def test_opdef_func():
+    sum_mt = mt.Sum([[1, 2]], [1])
+    sum_tf = sum_mt.reify()
+
+    with tf.compat.v1.Session() as sess:
+        assert sum_tf.eval() == np.r_[3]