Skip to content

Pull requests: google-deepmind/optax

New pull request New
77 Open 1,221 Closed
77 Open 1,221 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(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 AttributeError in Sharpness-Aware Minimization (SAM) example notebook by using _replace for NamedTuple
#1742 opened Jul 28, 2026 by rajasekharporeddy Collaborator Loading…
Fixes a TypeError in rosenbrock_ademamix.ipynb caused by a parameter name mismatch.
#1741 opened Jul 28, 2026 by rajasekharporeddy Collaborator Loading…
Fix the Attribute Error in gradient_accumulation_and_microbatching.ipynb example notebook
#1740 opened Jul 28, 2026 by rajasekharporeddy Collaborator Loading…
Fix the Attribute Error in Differentially private CNN on MNIST example notebook
#1739 opened Jul 28, 2026 by rajasekharporeddy Collaborator 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
1
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…
Fix projection_box crash on scalar bounds with multi-leaf pytrees
#1724 opened Jul 12, 2026 by winklemad Loading…
add missing docstring to get_spectral_radius_upper_bound in linear_algebra.py
#1721 opened Jul 4, 2026 by RavSinghChandan Loading…
4
add missing docstrings to tile_second_to_last_dim, multi_normal and parse_version
#1720 opened Jul 4, 2026 by RavSinghChandan Loading…
4
contrib: add SCION optimizer (norm-constrained LMO generalisation of Muon)
#1718 opened Jun 27, 2026 by irhyl Contributor Loading…
1
ProTip! Add no:assignee to see everything that’s not assigned.