An implementation of the new schedule-free SGD and AdamW optimisers in JAX using Optimistix.
To install, simply download this repo and run
pip install -e .
in your Python environment. Then use from schedule_free_optx import ScheduleFreeSGD, ScheduleFreeAdamW and use them as any other Optimistix optimiser!