message: Add additional tests for Mirrored and ReadOnSync variables.
  Adds test for ONLY_FIRST_REPLICA aggregation for Mirrored variables.  Adds tests for setting and modifying per-replica values for SyncOnRead variables before reading/syncing them.
diff: @@ -680,6 +680,35 @@ class MirroredVariableTest(test.TestCase, parameterized.TestCase):
  foo()
  + @combinations.generate(
  + combinations.combine(
  + distribution=[
  + strategy_combinations.mirrored_strategy_with_one_cpu,
  + strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
  + strategy_combinations.tpu_strategy,
  + strategy_combinations.central_storage_strategy_with_two_gpus,
  + ],
  + mode=["graph", "eager"]))
  + def testAggregationOnlyFirstReplica(self, distribution):
  + with distribution.scope():
  + v = variable_scope.variable(
  + 15.,
  + synchronization=variables_lib.VariableSynchronization.ON_WRITE,
  + aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA)
  + self.evaluate(variables_lib.global_variables_initializer())
  +
  + @def_function.function
  + def assign():
  + ctx = distribution_strategy_context.get_replica_context()
  + replica_id = ctx.replica_id_in_sync_group
  + return v.assign(math_ops.cast(replica_id, dtypes.float32))
  + per_replica_results = self.evaluate(distribution.experimental_local_results(
  + distribution.experimental_run_v2(assign)))
  + # The per-replica values should always match the first replicas value.
  + self.assertAllEqual(
  + array_ops.zeros(distribution.num_replicas_in_sync, dtypes.float32),
  + per_replica_results)
  +
  _TPU_STRATEGIES = (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1)
  @@ -757,6 +786,8 @@ def mirrored_and_tpu_strategy_combinations():
  mode=["graph", "eager"])
  +# TODO(b/144432582): Add variable aggregation type to combinations to simplify
  +# tests.
  def strategy_and_run_tf_function_combinations():
  # Test the combination of different strategies and whether a tf.function
  # is passed into strategy.experimental_run_v2."""
  @@ -1153,6 +1184,38 @@ class SyncOnReadVariableTest(test.TestCase, parameterized.TestCase):
  1.0 * i)
  self.assertEqual(per_replica_results, tuple(expected_result))
  + @combinations.generate(strategy_and_run_tf_function_combinations())
  + def testAssignPerReplicaBeforeRead(self, distribution,
  + experimental_run_tf_function):
  + aggregations = [
  + variables_lib.VariableAggregation.SUM,
  + variables_lib.VariableAggregation.MEAN,
  + variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
  + ]
  + for aggregation in aggregations:
  + with distribution.scope():
  + v = variable_scope.variable(
  + 0.,
  + synchronization=variables_lib.VariableSynchronization.ON_READ,
  + aggregation=aggregation)
  + self.evaluate(variables_lib.global_variables_initializer())
  +
  + def assign(var=v):
  + ctx = distribution_strategy_context.get_replica_context()
  + replica_id = ctx.replica_id_in_sync_group
  + return var.assign(math_ops.cast(replica_id, dtypes.float32))
  +
  + if experimental_run_tf_function:
  + assign = def_function.function(assign)
  +
  + per_replica_results = self.evaluate(
  + distribution.experimental_local_results(
  + distribution.experimental_run_v2(assign)))
  + expected_result = []
  + for i in range(distribution.num_replicas_in_sync):
  + expected_result.append(1.0 * i)
  + self.assertEqual(per_replica_results, tuple(expected_result))
  +
  @combinations.generate(mirrored_and_tpu_strategy_combinations())
  def testReadValueWithAggregationNoneInCrossReplicaContext(self, distribution):
  with distribution.scope():
  