Function test_scheduler_init

tests/optim/test_scheduler.py:296–323

(
    scheduler: ComposerScheduler,
    ssr: float,
    test_times: list[str],
    expected_lrs: list[float],
    dummy_schedulers_state: State,
)

    ],
)
def test_scheduler_init(
    scheduler: ComposerScheduler,
    ssr: float,
    test_times: list[str],
    expected_lrs: list[float],
    dummy_schedulers_state: State,
):

    state = dummy_schedulers_state
    assert state.dataloader_len is not None
    assert state.max_duration is not None
    state.max_duration = Time(value=int(state.max_duration.value * ssr), unit=state.max_duration.unit)
    for test_time, expected_lr in zip(test_times, expected_lrs):
        parsed_time = Time.from_timestring(test_time)
        assert parsed_time.unit in [TimeUnit.EPOCH, TimeUnit.BATCH]
        if parsed_time.unit == TimeUnit.EPOCH:
            state.timestamp = state.timestamp.copy(
                epoch=parsed_time,
                batch=Time(int(state.dataloader_len) * int(parsed_time), TimeUnit.BATCH),
            )
        else:
            state.timestamp = state.timestamp.copy(
                batch=parsed_time,
                epoch=Time(int(parsed_time) // int(state.dataloader_len), TimeUnit.EPOCH),
            )

        lr = scheduler(state, ssr)
        assert lr == pytest.approx(expected_lr, abs=1e-3)


@pytest.mark.parametrize(