Fix unsuccessful opdef lookup when function was not created via Defun.

tensorflower-gardener · tensorflower-gardener · commit eb281da2d437 · 2017-03-10T21:04:30.000-08:00
Change: 149826365
diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py
@@ -68,12 +68,7 @@ def _get_node_def(op):
 
 
 def _get_op_def(op):
-  # pylint: disable=protected-access
-  if hasattr(op, "_sig"):
-    return getattr(op, "_sig")
-  else:
-    return op_def_registry.get_registered_ops()[op.type]
-  # pylint: enable=protected-access
+  return op.op_def or op_def_registry.get_registered_ops()[op.type]
 
 
 def _is_in_placeholders(op, func_arg_placeholders):
@@ -248,8 +243,8 @@ def _call(sig, *inputs, **kwargs):
         output_types,
         name=name,
         attrs=attrs,
+        op_def=sig,
         compute_shapes=False)
-  setattr(op, "_sig", sig)  # Remember the signature.
   if op.outputs:
     if len(op.outputs) == 1:
       ret = op.outputs[0]
@@ -452,6 +447,7 @@ def __init__(self,
     self._shape_func = shape_func
     self._extra_kwargs = kwargs
     self._definition = None  # Constructed lazily.
+    self._sub_functions = dict()  # Constructed with definition.
 
     self._args = []
     assert isinstance(input_types, (list, tuple))
diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py
@@ -526,6 +526,48 @@ def BN1(x):
       self.assertAllClose(vals[0], vals[1])
       self.assertAllClose(vals[2], vals[3])
 
+  def testDeclare(self):
+    foo = function.Declare("Foo", [("x", dtypes.float32)],
+                           [("y", dtypes.float32)])
+
+    @function.Defun(dtypes.float32, func_name="Foo", out_names=["y"])
+    def FooImpl(x):
+      return x * x + 1
+
+    x = array_ops.placeholder(dtypes.float32)
+    y = foo(x)
+
+    g = ops.get_default_graph()
+    FooImpl.add_to_graph(g)
+
+    with self.test_session():
+      rand = np.random.uniform(size=(3, 3))
+      expected = rand * rand + 1.0
+      self.assertAllClose(expected, y.eval(feed_dict={x: rand}))
+
+  def testDeclareUsedInDefun(self):
+    foo = function.Declare("Foo", [("x", dtypes.float32)],
+                           [("y", dtypes.float32)])
+
+    @function.Defun()
+    def Bar(x):
+      return foo(x)
+
+    @function.Defun(dtypes.float32, func_name="Foo", out_names=["y"])
+    def FooImpl(x):
+      return x * x + 1
+
+    x = array_ops.placeholder(dtypes.float32)
+    y = Bar(x)
+
+    g = ops.get_default_graph()
+    FooImpl.add_to_graph(g)
+
+    with self.test_session():
+      rand = np.random.uniform(size=(3, 3))
+      expected = rand * rand + 1.0
+      self.assertAllClose(expected, y.eval(feed_dict={x: rand}))
+
   def testDeclareTypeMistake(self):
     foo = function.Declare("Foo", [("x", dtypes.float32)],
                            [("y", dtypes.float32)])
diff --git a/tensorflow/python/framework/importer_test.py b/tensorflow/python/framework/importer_test.py
@@ -981,6 +981,29 @@ def InnerFunc(x):
         self.assertEqual(sess.run("external:0"), 11)
         self.assertEqual(sess.run("outer:0"), 21)
 
+  def testImportInsideDefun(self):
+    g = ops.Graph()
+    with g.as_default():
+      @function.Defun()
+      def Add2(x, y):
+        return math_ops.add(x, y)
+
+      x = constant_op.constant(3.0, dtype=dtypes.float32)
+      y = constant_op.constant(-5.0, dtype=dtypes.float32)
+      z = Add2(x, y, name="z")  # pylint: disable=unexpected-keyword-arg
+
+    gdef = g.as_graph_def()
+
+    @function.Defun()
+    def TestFunc():
+      return importer.import_graph_def(gdef, return_elements=["z:0"])[0]
+
+    z = TestFunc()
+
+    with self.test_session():
+      z_val = z.eval()
+      self.assertEqual(z_val, -2.0)
+
 
 if __name__ == "__main__":
   test.main()
