def constant_schedule_with_cosine_burnin(
    init_step_size: float, final_step_size: float, burnin_steps: int
) -> StepSchedule:
    """Create a constant schedule with cosine burn-in.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    final_step_size: float
        The desired final step size.
    burnin_steps: int
        The length of burn-in, in steps.

    Returns
    -------
    schedule_fn: StepSchedule
    """
    if not init_step_size >= 0:
        raise ValueError("`init_step_size` should be >= 0.")
    if not final_step_size >= 0:
        raise ValueError("`final_step_size` should be >= 0.")
    if not burnin_steps >= 0:
        raise ValueError("`burnin_steps` should be >= 0.")

    def schedule(step: Array):
        t = jnp.minimum(step / burnin_steps, 1.0)
        coef = (1 + jnp.cos(t * np.pi)) * 0.5
        return coef * init_step_size + (1 - coef) * final_step_size

    return schedule
Create a constant schedule with cosine burn-in.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    final_step_size: float
        The desired final step size.
    burnin_steps: int
        The length of burn-in, in steps.

    Returns
    -------
    schedule_fn: StepSchedule
    
fortuna/prob_model/posterior/sgmcmc/sgmcmc_step_schedule.py