Issues

• CTC loss and gradients become NaN for valid alignments with masked logits

  type:bug
  google-deepmind/optax#1840 · Filter by author salma2vec opened on Oct 28, 2026
  1

• 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

• [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
  5

• 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

  

• 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
  1

  

• 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

• New Example on creating a custom GradientTransformation

  documentationtype:feature
  google-deepmind/optax#652 · Filter by author SauravMaheshkar opened on Dec 2, 2023
  2

  

• Pytree-based Optimizers

  type:feature
  google-deepmind/optax#432 · Filter by author cgarciae opened on Oct 4, 2022
  3

• Problems when jitting Adafactor with inject_hyperparams.

  type:bug
  google-deepmind/optax#412 · Filter by author mkunesch opened on Sep 7, 2022
  

• Type errors for general pytrees

  type:support
  google-deepmind/optax#384 · Filter by author brentyi opened on Jul 31, 2022
  2