schedule = optax.piecewise_constant_schedule(
    init_value=0.001,
    boundaries_and_scales={
        5000: 0.1,   # At step 5000, multiply LR by 0.1
        8000: 0.1,   # At step 8000, multiply by 0.1 again
    }
)

        __
        __