Issues

• Zero-patience scheduling reduces learning rates on improvements and ignores plateaus

  type:bug
  google-deepmind/optax#1805 · Filter by author sylvesterkaczmarek opened on Oct 6, 2026

• Masked nonfinite values contaminate cosine-loss gradients

  type:bug
  google-deepmind/optax#1800 · Filter by author sylvesterkaczmarek opened on Oct 3, 2026

• AdaDelta, Adan, and NovoGrad produce NaN on float16 zero-gradient steps

  type:bug
  google-deepmind/optax#1789 · Filter by author Anu27n opened on Sep 18, 2026
  3

• RMS transforms produce NaN on float16 zero-gradient steps

  type:bug
  google-deepmind/optax#1784 · Filter by author vishesh9131 opened on Sep 15, 2026
  2

• [Feature Request] Add NEAT conflict-aware gradient transformation to optax.contrib

  type:feature
  google-deepmind/optax#1770 · Filter by author ItCodinTime opened on Sep 5, 2026
  1

• Adam-family optimizers (adam, amsgrad, adabelief, yogi) produce NaN on float16 zero-gradient steps

  type:bug
  google-deepmind/optax#1754 · Filter by author kartik1pandey opened on Aug 16, 2026

• LBFGS - Add functionality to skip updates

  type:feature
  google-deepmind/optax#1708 · Filter by author inversecrime opened on Jun 17, 2026
  4

• GramMuon and Sharding for Muon

  type:feature
  google-deepmind/optax#1660 · Filter by author pillow37 opened on Apr 14, 2026
  1

• Add Self-Scaled Broyden (SSBroyden) and SSBFGS optimizers

  type:feature
  google-deepmind/optax#1658 · Filter by author armbrusl opened on Apr 9, 2026
  2

• Add PolarExpress Variant to Muon

  type:feature
  google-deepmind/optax#1621 · Filter by author MarcMachaczek opened on Mar 6, 2026

• Feature request: Adahessian optimizer

  type:feature
  google-deepmind/optax#1603 · Filter by author viralvgupta opened on Feb 17, 2026
  6

• Implement SOAP optimizer

  type:feature
  google-deepmind/optax#1580 · Filter by author Rahul-2k4 opened on Jan 25, 2026
  5

• Implementing AdaMuon

  type:feature
  google-deepmind/optax#1575 · Filter by author raghulchandramouli opened on Jan 24, 2026
  4

• [Feature Request] Adding the MARS Optimizer (Variance Reduction) from Hu et al. 2024

  type:feature
  google-deepmind/optax#1561 · Filter by author KumarADITHYA123 opened on Jan 13, 2026
  1

• Feature Request: add MixMatch and FixMatch training support for semi-supervised learning (SSL)

  type:feature
  google-deepmind/optax#1550 · Filter by author surajyadav-research opened on Jan 6, 2026
  3

• Feature request: add BYOL, SimSiam, DINO and Barlow Twins losses to optax.losses

  type:feature
  google-deepmind/optax#1528 · Filter by author surajyadav-research opened on Dec 10, 2025
  20

• Code for DoG is duplicated

  type:others
  google-deepmind/optax#1509 · Filter by author carlosgmartin opened on Nov 23, 2025
  6

  

• Publishing with py.typed

  type:feature
  google-deepmind/optax#1388 · Filter by author tmke8 opened on Jul 21, 2025
  3

  

• Add pre-commit hook that strips notebooks of metadata

  dependenciestype:feature
  google-deepmind/optax#1372 · Filter by author carlosgmartin opened on Jul 6, 2025
  3

• Minor suggestions for make_perturbed_fun

  type:feature
  google-deepmind/optax#1342 · Filter by author carlosgmartin opened on Jun 19, 2025
  4

• Add missing projections from jaxopt

  type:feature
  google-deepmind/optax#1280 · Filter by author carlosgmartin opened on Apr 27, 2025
  14

• Rename optax.perturbations distribution method log_prob to unnormalized_log_prob

  type:feature
  google-deepmind/optax#1214 · Filter by author carlosgmartin opened on Mar 1, 2025
  8

• Updates dtype do not need to match params dtype but only grads dtype a priori

  type:feature
  google-deepmind/optax#1098 · Filter by author vroulet opened on Oct 9, 2024
  

• Support for loss function with auxiliary data in linesearch

  type:feature
  google-deepmind/optax#1053 · Filter by author ro0mquy opened on Sep 10, 2024
  4

• Timeline for JaxOpt migration

  type:others
  google-deepmind/optax#977 · Filter by author Joshuaalbert opened on May 28, 2024
  35