# optax.adam(learning_rate) is roughly equivalent to:
optax.chain(
    optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8),
    optax.scale(-learning_rate)
)

        __
        __