message: Add functionality to normalization layer to work in the cross-replica context
diff: @@ -414,10 +414,19 @@ class BatchNormalizationV2(Layer):
  def _assign_moving_average(self, variable, value, momentum):
  with ops.name_scope(None, 'AssignMovingAvg',
  [variable, value, momentum]) as scope:
  + # TODO(b/120571621): We want to avoid colocating the variables here
  + # since TPUStrategy does not implement replica local variables.
  + # Remove this hack once we support TPULocalVariables.
  + is_tpu_strategy = False
  + if distribution_strategy_context.has_distribution_strategy():
  + distribute = distribution_strategy_context.get_distribution_strategy()
  + if distribute.__class__.__name__ == 'TPUStrategy':
  + is_tpu_strategy = True
  +
  # TODO(apassos,srbs,skyewm): the colocation constraints here are disabled
  # because of a bug which leads cond_v2 to skip rewriting them creating
  # conflicts.
  - if tf2.enabled():
  + if tf2.enabled() or is_tpu_strategy:
  cm = contextlib.contextmanager(lambda: (yield))()
  else:
  cm = ops.colocate_with(variable)
  @@ -655,12 +664,32 @@ class BatchNormalizationV2(Layer):
  d = _broadcast(array_ops.stop_gradient(d, name='renorm_d'))
  scale, offset = _compose_transforms(r, d, scale, offset)
  + if distribution_strategy_context.in_cross_replica_context():
  + strategy = distribution_strategy_context.get_distribution_strategy()
  def _do_update(var, value):
  + """Compute the updates for mean and variance."""
  + if in_eager_mode and not self.trainable:
  + return
  + return strategy.extended.update(
  + var, self._assign_moving_average, (value, self.momentum),
  + group=False)
  + # We need to unwrap the moving_mean or moving_variance in the case of
  + # training being false to match the output of true_fn and false_fn
  + # in the smart cond.
  + mean_update = tf_utils.smart_cond(
  + training,
  + lambda: _do_update(self.moving_mean, new_mean),
  + lambda: strategy.unwrap(self.moving_mean))
  + variance_update = tf_utils.smart_cond(
  + training,
  + lambda: _do_update(self.moving_variance, new_variance),
  + lambda: strategy.unwrap(self.moving_variance))
  + else:
  + def _do_update(var, value):
  + """Compute the updates for mean and variance."""
  if in_eager_mode and not self.trainable:
  return
  -
  return self._assign_moving_average(var, value, self.momentum)
  -
  mean_update = tf_utils.smart_cond(
  training,
  lambda: _do_update(self.moving_mean, new_mean),
  