schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,           # Start at 0
    peak_value=0.001,         # Warm up to this
    warmup_steps=1000,        # Over this many steps
    decay_steps=9000,         # Then decay over this many steps
    end_value=0.0001          # Down to this final value
)

        __
        __