def stepped_schedule(
    init_value: float, boundaries_and_scales: Optional[Dict[int, float]] = None
) -> base.Schedule:
    """
    Returns a stepped learning rate schedule with optional warmup.

    A stepped learning rate schedule decreases the learning rate
    by specified amounts at specified epochs.
    A common ImageNet schedule decays the learning rate by a factor of 0.1 at epochs 30, 60 and 80.
    This would be specified setting `boundaries_and_scales` to:

      {
        30: 0.1,
        60: 0.01,
        80: 0.001
      }

    Parameters
    ----------
    init_value: float
        An initial value `init_v`.
    boundaries_and_scales: Optional[Dict[int, float]]
        The schedule as a dict like `{epoch_x: lr_factor_x, epoch_y: lr_factor_y}`;
    the step occurs at epoch `epoch` and sets the learning rate to `init_value * lr_factor`

    Returns
    -------
    base.Schedule
        A function that maps step counts to values.
    """
    if boundaries_and_scales is not None:
        all_positive = all(scale >= 0.0 for scale in boundaries_and_scales.values())
        if not all_positive:
            raise ValueError(
                "`piecewise_constant_schedule` expects non-negative scale factors"
            )

    def schedule(count):
        v = init_value
        if boundaries_and_scales is not None:
            for threshold, scale in sorted(boundaries_and_scales.items()):
                indicator = jnp.maximum(0.0, jnp.sign(threshold - count))
                v = v * indicator + (1 - indicator) * scale * init_value
        return v

    return schedule

    Returns a stepped learning rate schedule with optional warmup.

    A stepped learning rate schedule decreases the learning rate
    by specified amounts at specified epochs.
    A common ImageNet schedule decays the learning rate by a factor of 0.1 at epochs 30, 60 and 80.
    This would be specified setting `boundaries_and_scales` to:

      {
        30: 0.1,
        60: 0.01,
        80: 0.001
      }

    Parameters
    ----------
    init_value: float
        An initial value `init_v`.
    boundaries_and_scales: Optional[Dict[int, float]]
        The schedule as a dict like `{epoch_x: lr_factor_x, epoch_y: lr_factor_y}`;
    the step occurs at epoch `epoch` and sets the learning rate to `init_value * lr_factor`

    Returns
    -------
    base.Schedule
        A function that maps step counts to values.
    
fortuna/utils/optimizer.py