code: def mtr_lm_dense(sz):
    
    n = 2 ** sz
    hparams = mtf_unitransformer_base()
    hparams.d_model = 1024
    hparams.max_length = 1024
    hparams.batch_size = 128
    # Parameters for my_layer_stack()
    hparams.num_hidden_layers = 6
    hparams.d_ff = 8192 * n
    hparams.d_kv = 256
    hparams.num_heads = 8 * n
    hparams.learning_rate_decay_steps = 65536
    hparams.layout = "batch:batch;vocab:model;d_ff:model;heads:model"
    hparams.mesh_shape = "batch:32"
    return hparams
docs: Series of architectures for language modeling.
  
    We assume infinite training data, so no dropout necessary.
  
    You can use languagemodel_wiki_noref_v32k_l1k.
    (1 epoch = ~46000 steps).
    TODO(noam): find a large enough dataset for these experiments.
  
    Args:
      sz: an integer
  
    Returns:
      a hparams