code: def mtf_transformer_lm_baseline():
    
    hparams = mtf_transformer_paper_lm(-1)
    hparams.batch_size = 128
    hparams.learning_rate_decay_steps = 27200  # one epoch on lm1b
    hparams.mesh_shape = "batch:8"
    return hparams
docs: Small language model to run on 1 TPU.
  
    Run this on 2x2 on languagemodel_lm1b32k_packed for 272000 steps (10 epochs)
    Results:
           params/10^9  log-ppl(per-token)
           0.14         3.202
  
    Returns:
      a hparams