code: def transformer_tall_pretrain_lm_tpu():
    
    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
docs: Hparams for transformer on LM pretraining on TPU with AdamW.