import copy


def set_global_layer_norm_config(
    base_config,
    layer_norm_config):
  """Sets all layers' softmax configuration in `base_config` to `layer_norm_config`."""
  def set_layer_norm(cfg):
    # We unlock the configs since the 'layer_norm' field might not already
    # exist.
    with cfg.unlocked():
      cfg.layer_norm = copy.deepcopy(layer_norm_config)

  base_config = copy.deepcopy(base_config)
  set_layer_norm(base_config.model_hparams.encoder)
  set_layer_norm(base_config.model_hparams.decoder)
  for block in base_config.model_hparams.encoder.encoder_1d_blocks:
    set_layer_norm(block)
  for block in base_config.model_hparams.decoder.encoder_decoder_1d_blocks:
    set_layer_norm(block)
  return base_config