Skip to content

Pull requests: google-deepmind/optax

New pull request New
80 Open 1,229 Closed
80 Open 1,229 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 masking in label smoothing
#1760 opened Aug 20, 2026 by tandede Loading…
docs: fix typos and update docstring formatting across multiple modules
#1759 opened Aug 19, 2026 by rajasekharporeddy Collaborator Loading…
Add missing pyrefly suppressions to fix main CI
#1757 opened Aug 16, 2026 by kartik1pandey Loading…
1
Fix NaN in scale_by_adam on float16 zero-gradient steps
#1755 opened Aug 16, 2026 by kartik1pandey Loading…
1
2
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…
5
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…
docs: fix COCOB docstring reference and author name
#1737 opened Jul 27, 2026 by wahid18-maqs Loading…
Support scalar parameters in scale_by_sm3
#1736 opened Jul 26, 2026 by TheSaiEaranti Loading…
Compute exact matrix powers in matrix_inverse_pth_root for any p
#1735 opened Jul 26, 2026 by TheSaiEaranti Loading…
Allow microbatching functions (optax.microbatching) to be called with dynamic shape dimensions or under jax.eval_shape without raising unhashable Tracer errors.
#1732 opened Jul 22, 2026 by copybara-service Bot Loading…
1
docs(contrib): fix paper reference link for ADOPT optimizer
#1731 opened Jul 22, 2026 by himoto Loading…
1
Treat integer hyperparameters as static in inject_hyperparams (fixes #412)
#1730 opened Jul 21, 2026 by ArneshBanerjee Loading…
1
2
Add nbstripout pre-commit hook, keeping notebook outputs (#1372)
#1729 opened Jul 20, 2026 by shubhangithub Loading…
1
1
Use jax.Array instead of jax.typing.ArrayLike across the codebase.
#1728 opened Jul 14, 2026 by copybara-service Bot Loading…
Add optional batch shape argument to tree_random_like.
#1727 opened Jul 13, 2026 by carlosgmartin Contributor Loading…
Informational, do not merge
#1726 opened Jul 13, 2026 by copybara-service Bot Loading…
1
Fix scale_by_rprop returning a one-step-stale update
#1725 opened Jul 12, 2026 by winklemad Loading…
ProTip! Type g i on any issue or pull request to go back to the issue listing page.