[tab] utils/lr_utils.py (python)
import optax
import numpy as np


def get_lr_schedule(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    if lr_schedule == "cos":
        assert (
            warmup_steps <= total_steps
        ), "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end,
        )
    elif lr_schedule == "wsd":
        assert (
            warmup_steps + wsd_decay_steps <= total_steps
        ), "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(
                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
            ),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(
                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
            ),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(
            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
        )


def get_host_lr_value_fn(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
):
    """Return a host-only callable that computes the LR for a given integer step.

    Mirrors the shapes of schedules created in get_lr_schedule but avoids JAX entirely,
    preventing host↔device transfers when used for logging/metrics.
    """

    def clip01(x: float) -> float:
        return float(min(1.0, max(0.0, x)))

    if lr_schedule == "cos":
        # Matches optax.warmup_cosine_decay_schedule semantics.
        def value(step: int) -> float:
            s = int(step)
            if warmup_steps > 0 and s < warmup_steps:
                progress = s / float(warmup_steps)
                return init_lr + (max_lr - init_lr) * progress
            # After warmup: cosine decay from peak (max_lr) to end_value over [warmup_steps, total_steps]
            denom = max(1, total_steps - warmup_steps)
            progress = clip01((s - warmup_steps) / float(denom))
            cosine_decay = 0.5 * (1.0 + float(np.cos(np.pi * progress)))
            return decay_end + (max_lr - decay_end) * cosine_decay

        return value

    if lr_schedule == "wsd":
        # Warmup → Sustain → Decay schedule joined at boundaries
        def value(step: int) -> float:
            s = int(step)
            if warmup_steps > 0 and s < warmup_steps:
                progress = s / float(warmup_steps)
                return init_lr + (max_lr - init_lr) * progress
            sustain_end = max(0, total_steps - wsd_decay_steps)
            if s < sustain_end:
                return max_lr
            # Linear decay from max_lr → decay_end over wsd_decay_steps
            denom = max(1, wsd_decay_steps)
            progress = clip01((s - sustain_end) / float(denom))
            return max_lr + (decay_end - max_lr) * progress

        return value

    raise ValueError(f"Learning rate schedule not supported: {lr_schedule}")