def cosine_schedule(init_step_size: float, total_steps: int) -> StepSchedule:
    """Create a cosine step schedule.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    total_steps: int
        The cycle length, in steps.

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

    def schedule(step: Array):
        t = step / total_steps
        return 0.5 * init_step_size * (1 + jnp.cos(t * np.pi))

    return schedule
Create a cosine step schedule.

    Parameters
    ----------
    init_step_size: float
        The initial step size.
    total_steps: int
        The cycle length, in steps.

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