message: Factor out the all-reduce in Optimizer to its own method.
diff: @@ -443,10 +443,14 @@ class OptimizerV2(trackable.Trackable):
  args=(grads_and_vars,),
  kwargs={"name": name})
  + def _aggregate_gradients(self, distribution, grads_and_vars):
  + """Returns all-reduced gradients."""
  + return distribution.extended.batch_reduce_to(
  + ds_reduce_util.ReduceOp.SUM, grads_and_vars)
  +
  def _distributed_apply(self, distribution, grads_and_vars, name, apply_state):
  """`apply_gradients` using a `DistributionStrategy`."""
  - reduced_grads = distribution.extended.batch_reduce_to(
  - ds_reduce_util.ReduceOp.SUM, grads_and_vars)
  + reduced_grads = self._aggregate_gradients(distribution, grads_and_vars)
  var_list = [v for _, v in grads_and_vars]
  grads_and_vars = zip(reduced_grads, var_list)
  