from .adamw import (
    get_adamw_with_cosine_scheduler,
    get_adamw_with_warm_up_cosine_scheduler,
    get_adamw_with_warmup_linear_scheduler,
    get_adamw_with_linear_scheduler
)
from .lion import (
    get_lion_with_cosine_scheduler,
    get_lion_with_with_warmup_linear_scheduler,
    get_lion_with_warm_up_cosine_scheduler,
    get_lion_with_linear_scheduler
)
from .adafactor import (
    get_adafactor_with_cosine_scheduler,
    get_adafactor_with_warm_up_cosine_scheduler,
    get_adafactor_with_warmup_linear_scheduler,
    get_adafactor_with_linear_scheduler
)

from .rmsprop import (
    get_rmsprop_with_cosine_scheduler,
    get_rmsprop_with_linear_scheduler,
    get_rmsprop_with_warmup_linear_scheduler,
    get_rmsprop_with_warm_up_cosine_scheduler
)

from .optimizer_utils import optax_add_scheduled_weight_decay
