message: Do nothing in `add_update` when in cross-replica mode.
diff: @@ -774,11 +774,6 @@ class Layer(module.Module):
  with backend.get_graph().as_default():
  updates = []
  for u in self._updates:
  - # Filter out updates created in a cross-replica context when in a
  - # replica context and vice versa.
  - if (getattr(u, '_in_cross_replica_context', False) !=
  - ds_context.in_cross_replica_context()):
  - continue
  if callable(u):
  try:
  u = u()
  @@ -1059,6 +1054,15 @@ class Layer(module.Module):
  on this Layer, when executing in Eager mode.
  inputs: Deprecated, will be automatically inferred.
  """
  + if ds_context.has_strategy() and ds_context.in_cross_replica_context():
  + # Updates don't need to be run in a cross-replica context.
  + if (ops.executing_eagerly_outside_functions() and
  + not base_layer_utils.is_in_keras_graph()):
  + raise RuntimeError( # pylint: disable=g-doc-exception
  + '`add_update` was called in a cross-replica context. This is not '
  + 'expected. If you require this feature, please file an issue.')
  + return
  +
  updates = generic_utils.to_list(updates)
  call_context = base_layer_utils.call_context()
  @@ -1100,8 +1104,6 @@ class Layer(module.Module):
  reachable = tf_utils.get_reachable_from_inputs(relevant_inputs, [update])
  update._unconditional_update = update not in reachable
  - update._in_cross_replica_context = (
  - ds_context.has_strategy() and ds_context.in_cross_replica_context())
  return update
  updates = [process_update(x) for x in updates]
  