# optax.adamw(learning_rate, weight_decay) is roughly:
optax.chain(
    optax.scale_by_adam(),
    optax.add_decayed_weights(weight_decay),
    optax.scale(-learning_rate)
)

        __
        __