from __future__ import annotations

from typing import Any, Optional

class LRScheduler:
    base_lr: Any = ...
    last_lr: Any = ...
    last_epoch: Any = ...
    verbose: Any = ...
    def __init__(self, learning_rate: float = ..., last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def __call__(self): ...
    def step(self, epoch: Optional[Any] = ...) -> None: ...
    def state_dict(self): ...
    keys: Any = ...
    def state_keys(self) -> None: ...
    def set_state_dict(self, state_dict: Any) -> None: ...
    set_dict: Any = ...
    def get_lr(self) -> None: ...

class NoamDecay(LRScheduler):
    d_model: Any = ...
    warmup_steps: Any = ...
    def __init__(
        self, d_model: Any, warmup_steps: Any, learning_rate: float = ..., last_epoch: int = ..., verbose: bool = ...
    ) -> None: ...
    def get_lr(self): ...

class PiecewiseDecay(LRScheduler):
    boundaries: Any = ...
    values: Any = ...
    def __init__(self, boundaries: Any, values: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...

class NaturalExpDecay(LRScheduler):
    gamma: Any = ...
    def __init__(self, learning_rate: Any, gamma: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...

class InverseTimeDecay(LRScheduler):
    gamma: Any = ...
    def __init__(self, learning_rate: Any, gamma: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...

class PolynomialDecay(LRScheduler):
    decay_steps: Any = ...
    end_lr: Any = ...
    power: Any = ...
    cycle: Any = ...
    def __init__(
        self,
        learning_rate: Any,
        decay_steps: Any,
        end_lr: float = ...,
        power: float = ...,
        cycle: bool = ...,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def get_lr(self): ...

class LinearWarmup(LRScheduler):
    learning_rate: Any = ...
    warmup_steps: Any = ...
    start_lr: Any = ...
    end_lr: Any = ...
    def __init__(
        self,
        learning_rate: Any,
        warmup_steps: Any,
        start_lr: Any,
        end_lr: Any,
        last_epoch: int = ...,
        verbose: bool = ...,
    ) -> None: ...
    def state_dict(self): ...
    def set_state_dict(self, state_dict: Any) -> None: ...
    def get_lr(self): ...

class ExponentialDecay(LRScheduler):
    gamma: Any = ...
    def __init__(self, learning_rate: Any, gamma: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...

class MultiStepDecay(LRScheduler):
    milestones: Any = ...
    gamma: Any = ...
    def __init__(
        self, learning_rate: Any, milestones: Any, gamma: float = ..., last_epoch: int = ..., verbose: bool = ...
    ) -> None: ...
    def get_lr(self): ...

class StepDecay(LRScheduler):
    step_size: Any = ...
    gamma: Any = ...
    def __init__(
        self, learning_rate: Any, step_size: Any, gamma: float = ..., last_epoch: int = ..., verbose: bool = ...
    ) -> None: ...
    def get_lr(self): ...

class LambdaDecay(LRScheduler):
    lr_lambda: Any = ...
    def __init__(self, learning_rate: Any, lr_lambda: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...

class ReduceOnPlateau(LRScheduler):
    mode: Any = ...
    factor: Any = ...
    threshold_mode: Any = ...
    patience: Any = ...
    threshold: Any = ...
    cooldown: Any = ...
    min_lr: Any = ...
    epsilon: Any = ...
    cooldown_counter: int = ...
    best: Any = ...
    num_bad_epochs: int = ...
    base_lr: Any = ...
    last_lr: Any = ...
    last_epoch: int = ...
    verbose: Any = ...
    def __init__(
        self,
        learning_rate: Any,
        mode: str = ...,
        factor: float = ...,
        patience: int = ...,
        threshold: float = ...,
        threshold_mode: str = ...,
        cooldown: int = ...,
        min_lr: int = ...,
        epsilon: float = ...,
        verbose: bool = ...,
    ) -> None: ...
    keys: Any = ...
    def state_keys(self) -> None: ...
    def step(self, metrics: Any, epoch: Optional[Any] = ...) -> None: ...

class CosineAnnealingDecay(LRScheduler):
    T_max: Any = ...
    eta_min: Any = ...
    def __init__(
        self, learning_rate: Any, T_max: Any, eta_min: int = ..., last_epoch: int = ..., verbose: bool = ...
    ) -> None: ...
    def get_lr(self): ...

class MultiplicativeDecay(LRScheduler):
    lr_lambda: Any = ...
    def __init__(self, learning_rate: Any, lr_lambda: Any, last_epoch: int = ..., verbose: bool = ...) -> None: ...
    def get_lr(self): ...
