[DRAFT] Modify CosWithWarmupAndLinearDecay tail-handoff and missing super().__post_init__()

[AkshitaB]

Summary

TODO: unclear if we actually want this behavior.

Rewrites the tail-handoff semantics of CosWithWarmupAndLinearDecay and fixes a missing super().__post_init__() call. Original work by @Tianhua-Tao; this PR ports just the scheduler portion onto main (preserving authorship/date) and adds tests + a CHANGELOG entry.

What's changed

1. Tail starts at the cosine's terminal value, not its mid-flight value

Before: The cosine ran across the full t_max. At step t_max - decay, the linear tail attached at whatever value the cosine happened to produce there (call that cos(t_max - decay, t_max) — a value somewhere between peak and alpha_f * peak, depending on horizon and alpha_f). The tail then anneals from that mid-flight value to decay_min_lr.

After: The cosine completes (reaches alpha_f * peak) by step t_max - decay. Achieved by passing t_max - decay to super().get_lr(...) for the cosine region, and hard-coding the tail's starting value to initial_lr * alpha_f. The linear tail then anneals from alpha_f * peak down to decay_min_lr over the last decay steps.

Net effect: the class now matches the schedule its name promises — "warmup → cosine completing to alpha_f * peak by t_max - decay → linear tail to decay_min_lr over decay steps."

2. Missing super().__post_init__() call

CosWithWarmupAndLinearDecay.__post_init__ validated its own fields but never invoked super().__post_init__(). This meant the parent's deprecated-alias migration (warmup_steps → warmup) was skipped on the subclass, and the parent's warmup_fraction range check ran only via duplicated logic in the child. Now super().__post_init__() is called at the end.

Tests

src/test/optim/scheduler_test.py gains:

• test_cos_with_warmup_and_linear_decay — verifies the three-phase schedule. The cosine region matches a bare CosWithWarmup evaluated against t_max = (full_t_max - decay). The tail starts at alpha_f * peak at exactly t_max - decay, anneals linearly to decay_min_lr by t_max.
• test_cos_with_warmup_and_linear_decay_migrates_deprecated_warmup_steps — confirms the super().__post_init__() call now propagates the warmup_steps → warmup migration and emits the deprecation warning.

34 tests pass; make checks clean (isort, black, ruff, mypy across all 381 source files).

Commits

SHA	Author	Subject
0dd3f4c5	Tianhua Tao	Rewrite CosWithWarmupAndLinearDecay tail semantics (the core change)
f824210d	Akshita Bhagia	Clean up leftover # ... reference comments and trailing whitespace
b7ab9a43	Akshita Bhagia	Add tests
8834192c	Akshita Bhagia	Update CHANGELOG

Compatibility note

Existing callers of CosWithWarmupAndLinearDecay will see different LR values in the tail region. If you have runs in flight that depend on the old "tail attaches mid-cosine" behavior, this is a behavior change — but the class name already implies the new semantics, and the old behavior was buggy in the sense that the tail start was implicit/non-obvious.

Test plan

• pytest src/test/optim/scheduler_test.py — 34 passed
• make checks (isort + black + ruff + mypy) — clean
• Confirm no existing training scripts on main depend on the old tail-attachment behavior of CosWithWarmupAndLinearDecay

🤖 Generated with Claude Code