Hoist resource variable reads outside functions.

alextp · tensorflower-gardener · commit 020501c69511 · 2017-03-06T16:50:30.000-08:00
Change: 149360110
diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py
@@ -33,6 +33,7 @@
 from tensorflow.python.framework import op_def_registry
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.ops import variable_scope as vs
 from tensorflow.python.util import compat
 
@@ -324,6 +325,12 @@ def getvar(self,
           collections=collections,
           use_resource=use_resource)
       self.extra_vars.append(var)
+      if isinstance(var, resource_variable_ops.ResourceVariable):
+        # For resource-based variables read the variable outside the function
+        # and pass in the value. This ensures that the function is pure and
+        # differentiable. TODO(apassos) this may have performance problems if
+        # the function will only do embedding lookups on the variable.
+        return var.value()
       return var
 
   def create_op(self, op_type, inputs, data_types, **kwargs):
diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py
@@ -1178,9 +1178,7 @@ def testBasic(self):
     self._testSimpleModel(True)
     self._testSimpleModel(False)
 
-  # TODO(b/35668241): disabled because resource variable handling inside
-  # functions does not work.
-  def DISABLED_testBasicResource(self):
+  def testBasicResource(self):
     self._testSimpleModel(True, use_resource=True)
     self._testSimpleModel(False, use_resource=True)
 
