schedule free adamw jax update with switching params to x for validation, y for training#16

Open

wyfEmma wants to merge 11 commits intomlcommons:mainfrom

wyfEmma:wyf_schedule_free

Commits on Aug 24, 2025

Commits on Sep 1, 2025

adding schedule free adamw compatible with jit
init-22
committed

Commits on Feb 4, 2026

schedule free prepare_for_eval
wyfEmma
committed

Commits on Feb 22, 2026

Merge branch 'mlcommons:main' into wyf_schedule_free
wyfEmma
authored

Commits on Mar 1, 2026

1. move jitted train step declaration out of the update params function such to remove the unnecessary overhead. 2. fix imports schedule_free_eval_params 3. switch to jax.tree.map instead of jax.tr…
wyfEmma
committed
Merge branch 'wyf_schedule_free' of https://github.com/wyfEmma/submissions_algorithms into wyf_schedule_free
wyfEmma
committed

Commits on Mar 8, 2026

1. wrap schedule_free_eval_params function and restore y calculation in jit to reduce temporary memory peak
wyfEmma
committed

Commits on Mar 27, 2026

Change replication approach, use different
wyfEmma
committed

Commits on Mar 31, 2026