Add PolarExpress variant to Muon optimizer#1613

Open

MarcMachaczek wants to merge 20 commits into

google-deepmind:mainfrom

MarcMachaczek:feat/polar-express

Commits on Mar 4, 2026

Commits on Mar 5, 2026

relax tolerance for test
MarcMachaczek
committed
test mismatch reason solved: increased cutoff for l/u ratio for which remez algorithm is skipped from 1 - 5e-6 to 1 - 5e-5. for the previous cutoff, the linear system became ill-conditioned (k~1e11…
MarcMachaczek
committed
testing polar express with different preconditioners on hard matrices (binary and low-rank spectrum)
MarcMachaczek
committed
compare coefficients to hard coded coefficientw as found in the original work, this required a different cushion value
MarcMachaczek
committed
changing how safety factor is applied to the polynomial coefficients. this again allows for tighter remez algorithm cutoff (now again the same as in original work)
MarcMachaczek
committed
use default dtype independent lower bound for polar express of 1e-3
MarcMachaczek
committed
remove polar express preconditioning; use default (frobenius) and mention schatten as recommended; remove associated warnings ang logging import
MarcMachaczek
committed
removed polar express constants in favor of default arguments
MarcMachaczek
committed
final formatting changes to minimize diff with main
MarcMachaczek
committed

Commits on Mar 8, 2026