<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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