code: def afx_adam():
    
    hparams = transformer.transformer_base_v2()
    hparams.optimizer_adam_beta1 = 0.9
    hparams.optimizer_adam_beta2 = 0.999
    hparams.symbol_modality_num_shards = 1
    hparams.batch_size = 2048
    hparams.optimizer = "adam"
    hparams.learning_rate_schedule = (
        "constant*rsqrt_decay*linear_warmup*rsqrt_hidden_size")
    hparams.learning_rate_constant = 2.0
    return hparams
docs: Old version - Adam.