Eliminate the creation of unnecessary read ops when working with ResourceVariables.

In particular:
1. Don't create additional read ops when creating a ResourceVariable from a VariableDef proto.
2. Expose the ability to assign a ResourceVariable without reading & returning the new value.
3. Colocating with a ResourceVariable's ".op" property eliminates the creation of additional read ops.
4. Savers can read a variable's value using the _graph_element property, since these reads don't need control dependencies.

This makes the visualization of graphs on TensorBoard much nicer.

PiperOrigin-RevId: 187622122

1 files changed, 45 insertions, 0 deletions

diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
index 8503f3e031..71699fe0ad 100644
--- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py
+++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
@@ -277,6 +277,20 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase):
     self.evaluate(v.assign(2.0))
     self.assertEqual(2.0, self.evaluate(v.value()))
 
+    # Tests for the 'read_value' argument:
+    assign_with_read = v.assign(3.0, read_value=True)
+    if context.in_graph_mode():
+      self.assertEqual(3.0, assign_with_read.eval())
+    else:
+      self.assertEqual(3.0, self.evaluate(assign_with_read))
+    assign_without_read = v.assign(4.0, read_value=False)
+    if context.in_graph_mode():
+      self.assertIsInstance(assign_without_read, ops.Operation)
+    else:
+      self.assertIsNone(assign_without_read)
+    self.evaluate(assign_without_read)
+    self.assertEqual(4.0, self.evaluate(v.value()))
+
   @test_util.run_in_graph_and_eager_modes()
   def testLoad(self):
     v = resource_variable_ops.ResourceVariable(1.0, name="var0")
@@ -329,6 +343,9 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase):
       w = resource_variable_ops.ResourceVariable.from_proto(v.to_proto())
       self.assertEquals(2, math_ops.add(w, 1).eval())
 
+      self.assertEquals(v._handle, w._handle)
+      self.assertEquals(v._graph_element, w._graph_element)
+
   @test_util.run_in_graph_and_eager_modes()
   def testAssignAddMethod(self):
     v = resource_variable_ops.ResourceVariable(1.0, name="var0")
@@ -336,6 +353,20 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase):
     self.evaluate(v.assign_add(1.0))
     self.assertEqual(2.0, self.evaluate(v.value()))
 
+    # Tests for the 'read_value' argument:
+    assign_with_read = v.assign_add(1.0, read_value=True)
+    if context.in_graph_mode():
+      self.assertEqual(3.0, assign_with_read.eval())
+    else:
+      self.assertEqual(3.0, self.evaluate(assign_with_read))
+    assign_without_read = v.assign_add(1.0, read_value=False)
+    if context.in_graph_mode():
+      self.assertIsInstance(assign_without_read, ops.Operation)
+    else:
+      self.assertIsNone(assign_without_read)
+    self.evaluate(assign_without_read)
+    self.assertEqual(4.0, self.evaluate(v.value()))
+
   @test_util.run_in_graph_and_eager_modes()
   def testAssignSubMethod(self):
     v = resource_variable_ops.ResourceVariable(3.0, name="var0")
@@ -343,6 +374,20 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase):
     self.evaluate(v.assign_sub(1.0))
     self.assertEqual(2.0, self.evaluate(v.value()))
 
+    # Tests for the 'read_value' argument:
+    assign_with_read = v.assign_sub(1.0, read_value=True)
+    if context.in_graph_mode():
+      self.assertEqual(1.0, assign_with_read.eval())
+    else:
+      self.assertEqual(1.0, self.evaluate(assign_with_read))
+    assign_without_read = v.assign_sub(1.0, read_value=False)
+    if context.in_graph_mode():
+      self.assertIsInstance(assign_without_read, ops.Operation)
+    else:
+      self.assertIsNone(assign_without_read)
+    self.evaluate(assign_without_read)
+    self.assertEqual(0.0, self.evaluate(v.value()))
+
   @test_util.run_in_graph_and_eager_modes()
   def testDestroyResource(self):
     v = resource_variable_ops.ResourceVariable(3.0, name="var0")