code: def afx_small():
    
    hparams = transformer.transformer_tpu()
    hparams.filter_size = 1024
    hparams.num_heads = 4
    hparams.num_hidden_layers = 3
    hparams.batch_size = 512
    return hparams
docs: Small transformer model with small batch size for fast step times.