Function scale_from_learning_rate

axlearn/common/optimizers.py:271–284

(
    learning_rate: schedule.Schedule, *, flip_sign=True
)

def scale_from_learning_rate(
    learning_rate: schedule.Schedule, *, flip_sign=True
) -> schedule.ScheduleFn:
    learning_rate_fn = schedule.as_schedule_fn(learning_rate)

    def scale_fn(step):
        lr = learning_rate_fn(step)
        context = current_context()
        if context:
            context.add_summary("lr_schedule_step", step)
            context.add_summary("learning_rate", lr)
        return -lr if flip_sign else lr

    return scale_fn


def per_param_scale_by_path(