parameters: (opt)
docstring: Wrap the optimizer `opt` with CrossShardOptimizer if applicable.
docstring_summary: Wrap the optimizer `opt` with CrossShardOptimizer if applicable.
docstring_tokens: Wrap the optimizer opt with CrossShardOptimizer if applicable .
function: def _replicated_optimizer(opt):
    """Wrap the optimizer `opt` with CrossShardOptimizer if applicable."""
    # Always wrap `opt` with CrossShardOptimizer, even if we are running on a
    # single core.  This ensures Keras properly tracks and initializes optimizer
    # variables.
    if isinstance(opt, keras_optimizers.TFOptimizer):
      return tpu_optimizer.CrossShardOptimizer(opt.optimizer)
    else:
      return KerasCrossShardOptimizer(opt)
function_tokens: def _replicated_optimizer ( opt ) : # Always wrap `opt` with CrossShardOptimizer, even if we are running on a # single core.  This ensures Keras properly tracks and initializes optimizer # variables. if isinstance ( opt , keras_optimizers . TFOptimizer ) : return tpu_optimizer . CrossShardOptimizer ( opt . optimizer ) else : return KerasCrossShardOptimizer ( opt )