Skip to content

Pull requests: google-deepmind/optax

New pull request New
87 Open 1,236 Closed
87 Open 1,236 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Fix Pyrefly missing-attribute and unsupported-operation ignores on ArrayLike arguments across loss functions, transforms, linesearches, and schedules by normalizing them to jax.Array using jnp.asarray.
#1776 opened Sep 11, 2026 by copybara-service Bot Loading…
1
Relax tolerance in _muon_test for GPU execution.
#1775 opened Sep 10, 2026 by copybara-service Bot Loading…
Preserve plateau scale dtype with array minimum scales
#1774 opened Sep 10, 2026 by sylvesterkaczmarek Loading…
Fix Pyrefly missing-attribute and unsupported-operation ignores on ArrayLike values by normalizing them to jax.Array using jnp.asarray.
#1773 opened Sep 8, 2026 by copybara-service Bot Loading…
1
Fix negative CTC loss for confident alignments in float32
#1772 opened Sep 7, 2026 by rajasekharporeddy Collaborator Loading…
1
fix(transform): prevent NaN in Adam-family optimizers on float16 zero-gradient steps
#1769 opened Sep 3, 2026 by somuai Loading…
1
2
fix(transform): prevent NaN in scale_by_adamax on float16 zero-gradient steps
#1768 opened Sep 3, 2026 by somuai Loading…
2
Fix contrib.reduce_on_plateau never reducing the learning rate on negative values
#1767 opened Sep 2, 2026 by Kayvan-Zahiri Loading…
1
Fix a NaN in contrib.dadapt_adamw on zero denominators and zero learning rates.
#1766 opened Aug 30, 2026 by copybara-service Bot Loading…
Guard contrib.prodigy against zero denominators and zero learning rates.
#1765 opened Aug 30, 2026 by copybara-service Bot Loading…
fix(losses): avoid NaN gradient in weighted_logsoftmax at -inf logits
#1762 opened Aug 25, 2026 by jaideeppyne Loading…
Fix all-NaN output for complex inputs with 'schatten' preconditioning in Muon
#1761 opened Aug 24, 2026 by shoemoney Loading…
2
Fix masking in label smoothing
#1760 opened Aug 20, 2026 by tandede Loading…
2
docs: fix typos and update docstring formatting across multiple modules
#1759 opened Aug 19, 2026 by rajasekharporeddy Collaborator Loading…
1
Fix NaN in scale_by_adam on float16 zero-gradient steps
#1755 opened Aug 16, 2026 by kartik1pandey Loading…
1
7
Add SPSA (Simultaneous Perturbation Stochastic Approximation) optimizer to contrib
#1753 opened Aug 16, 2026 by jaideeppyne Loading…
1
2
Fix clip_by_global_norm producing NaN when the global norm and max_norm are both zero
#1752 opened Aug 15, 2026 by shashvat-singham Loading…
6
Support negative axes in microbatching
#1751 opened Aug 13, 2026 by anxkhn Loading…
2
Skip invalid L-BFGS curvature updates
#1749 opened Aug 6, 2026 by Trighap52 Loading…
1
1
fix(contrib): improve adopt docstring example to show meaningful convergence and add doctest coverage
#1747 opened Jul 31, 2026 by wahid18-maqs • Draft
fix(transforms): reset notfinite_count after apply_if_finite gives up and accepts non-finite update
#1746 opened Jul 31, 2026 by wahid18-maqs Loading…
fix(contrib): correct momo and momo_adam docstring examples to use apply_updates
#1745 opened Jul 31, 2026 by wahid18-maqs Loading…
fix: initialize schedule_free step_count to zero
#1744 opened Jul 31, 2026 by wahid18-maqs Loading…
fix: prevent division by zero in SAM gradient normalization
#1743 opened Jul 31, 2026 by wahid18-maqs Loading…
fix: warn when Hutchinson estimator uses fixed PRNGKey(0)
#1738 opened Jul 27, 2026 by wahid18-maqs Loading…
ProTip! Exclude everything labeled bug with -label:bug.