Skip to content

Pull requests: google-deepmind/optax

New pull request New
72 Open 1,221 Closed
72 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 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
Fix ACProp second moment to use the new first moment
#1716 opened Jun 27, 2026 by codewithfourtix Loading…
1
Fix incorrect (1 - b2) coefficient in scale_by_adan
#1715 opened Jun 26, 2026 by codewithfourtix Loading…
5
Add Linear Neuron Boosting to contrib
#1713 opened Jun 23, 2026 by dmunoz-research Loading…
feat: Add SPSA optimization method (Issue #357)
#1712 opened Jun 23, 2026 by Paramveersingh-S Loading…
7
Migrate Reduce on Plateau Learning Rate Scheduler example to the latest flax.nnx module
#1711 opened Jun 22, 2026 by rajasekharporeddy Collaborator Loading…
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.