mlcommons / mlcommons/algorithmic-efficiency

Add schedule-free adamw submission in JAX

Open
#809 20 comments 0 reactions 1 assignee View on GitHub

@init-22 is already working on this.

Since Aug 21, 2025.

⚡ Submission
Dominant language
Python
Stars
425
Forks
78
No merged PRs in 30d

Description

Description

Currently we have been unable to reproduce the schedule free adamw results with JAX.
There seem to be differences between the optax implementation of schedule-free adamw and the pytorch submission.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.