Hparams for transformer on LM pretraining on TPU with AdamW.

def transformer_tall_pretrain_lm_tpu():
  """Hparams for transformer on LM pretraining on TPU with AdamW."""
  hparams = transformer_tall_pretrain_lm_tpu_adafactor()
  # Optimizer gets reset in update_hparams_for_tpu so we set it again here.
  hparams.learning_rate_constant = 2e-4
  hparams.learning_rate_schedule = ("linear_warmup * constant * cosdecay")
  hparams.optimizer = "adam_w"
  return hparams