schedule = optax.exponential_decay(
    init_value=0.001,
    transition_steps=1000,
    decay_rate=0.96
)

        __
        __