code: def mtf_image_transformer_base_imagenet_mp128():
    
    hparams = mtf_image_transformer_base_imagenet()
    hparams.mesh_shape = "model:8;batch:4"
    hparams.layout = "batch:batch;d_ff:model;heads:model"
    hparams.batch_size = 8
    hparams.img_len = 128
    hparams.block_length = 128
    hparams.num_heads = 8
    hparams.num_decoder_layers = 4
    hparams.d_ff = 4096
    hparams.learning_rate_warmup_steps = 31250
    hparams.unconditional = True
    hparams.max_length = 256*256*3
    return hparams
docs: Model parallel ImageNet parameters.