def polynomial_schedule(
    a: float = 1.0, b: float = 1.0, gamma: float = 0.55
) -> StepSchedule:
    """Create a polynomial step schedule.

    Parameters
    ----------
    a: float
        Scale of all step sizes.
    b: float
        The stabilization constant.
    gamma: float
        The decay rate :math:`\gamma \in (0.5, 1.0]`.

    Returns
    -------
    schedule_fn: StepSchedule
    """

    if not 0.5 < gamma <= 1.0:
        raise ValueError("`gamma` should be in (0.5, 1.0] range.")

    def schedule(step: Array):
        return a * (b + step) ** (-gamma)

    return schedule
Create a polynomial step schedule.

    Parameters
    ----------
    a: float
        Scale of all step sizes.
    b: float
        The stabilization constant.
    gamma: float
        The decay rate :math:`\gamma \in (0.5, 1.0]`.

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