Fix calling tf.identity on SyncOnReadVariable in cross replica context.

rxsang · tensorflower-gardener · commit e78c47f2a73a · 2020-02-07T13:28:21.000-08:00
PiperOrigin-RevId: 293881729
Change-Id: I94bfcd4380767cdbbacf30ac5a29ee7bb2b992fa
diff --git a/tensorflow/python/distribute/values.py b/tensorflow/python/distribute/values.py
@@ -1139,8 +1139,9 @@ def _saveable_factory(name=self._common_name):
 
   def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
     """Converts a variable to a tensor."""
-    return ops.convert_to_tensor(
-        self._get(), dtype=dtype, name=name, as_ref=as_ref)
+    with _enter_or_assert_strategy(self._distribute_strategy):
+      return ops.convert_to_tensor(
+          self._get(), dtype=dtype, name=name, as_ref=as_ref)
 
 
 # Register a conversion function for SyncOnReadVariable which allows as_ref to
diff --git a/tensorflow/python/distribute/values_test.py b/tensorflow/python/distribute/values_test.py
@@ -867,19 +867,6 @@ def testProperties(self):
     self.assertEqual(variable_scope.VariableAggregation.SUM,
                      replica_local.aggregation)
 
-  def testTensorConversion(self):
-    with context.graph_mode():
-      _, replica_local = _make_replica_local(
-          variable_scope.VariableAggregation.SUM)
-      converted = ops.convert_to_tensor(replica_local, as_ref=False)
-      self.assertIsInstance(converted, ops.Tensor)
-      self.assertEqual(converted.dtype, replica_local.dtype)
-
-      converted = ops.convert_to_tensor(replica_local, as_ref=True)
-      # Resources variable are converted to tensors as well when as_ref is True.
-      self.assertIsInstance(converted, ops.Tensor)
-      self.assertEqual(converted.dtype, replica_local.dtype)
-
   @test_util.run_v2_only
   def testCanPassToDefFun(self):
     @def_function.function
@@ -937,6 +924,20 @@ def _save(self, sess, var):
     save_path, _ = self._save_return_saver(sess, var)
     return save_path
 
+  @combinations.generate(mirrored_and_tpu_strategy_combinations())
+  def testTensorConversion(self, distribution):
+    with context.graph_mode():
+      _, replica_local = _make_replica_local(
+          variable_scope.VariableAggregation.SUM, distribution)
+      converted = ops.convert_to_tensor(replica_local, as_ref=False)
+      self.assertIsInstance(converted, ops.Tensor)
+      self.assertEqual(converted.dtype, replica_local.dtype)
+
+      converted = ops.convert_to_tensor(replica_local, as_ref=True)
+      # Resources variable are converted to tensors as well when as_ref is True.
+      self.assertIsInstance(converted, ops.Tensor)
+      self.assertEqual(converted.dtype, replica_local.dtype)
+
   @combinations.generate(mirrored_and_tpu_strategy_combinations())
   def testSaveAndRestoreReplicaLocalSumOneGraph(self, distribution):
     with self.cached_session() as sess:
@@ -1240,7 +1241,7 @@ def testReadValueInCrossReplicaContext(self, distribution,
     ]
     for aggregation in aggregations:
       if isinstance(distribution, _TPU_STRATEGIES):
-        resolver = tpu_cluster_resolver.TPUClusterResolver('')
+        resolver = tpu_cluster_resolver.TPUClusterResolver("")
         tpu_strategy_util.initialize_tpu_system(resolver)
       with distribution.scope():
         v = variable_scope.variable(
@@ -1270,6 +1271,8 @@ def assign(v=v):
       self.assertEqual(expected, self.evaluate(v.read_value()), aggregation)
       self.assertEqual(expected, self.evaluate(v.value()), aggregation)
       self.assertEqual(expected, self.evaluate(v), aggregation)
+      self.assertEqual(expected, self.evaluate(array_ops.identity(v)),
+                       aggregation)
 
   # TODO(b/145574622): Re-enable this test once ReduceOp argument is
   # respected on GPUs.
