Skip to content

Pull requests: google-deepmind/optax

New pull requestNew
75 Open 1,226 Closed
75 Open 1,226 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 reviewsReview requiredApproved reviewChanges requested
Assignee
Filter by who’s assigned
Assigned to nobodyLoading
Sort
Sort by
NewestOldestMost commentedLeast commentedRecently updatedLeast recently updatedBest match
Most reactions
👍👎😄🎉😕❤️🚀👀

Pull requests list

Skip invalid L-BFGS curvature updates
#1749 opened Aug 6, 2026 by Trighap52Loading…
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-maqsLoading…
fix(contrib): correct momo and momo_adam docstring examples to use apply_updates
#1745 opened Jul 31, 2026 by wahid18-maqsLoading…
fix: initialize schedule_free step_count to zero
#1744 opened Jul 31, 2026 by wahid18-maqsLoading…
fix: prevent division by zero in SAM gradient normalization
#1743 opened Jul 31, 2026 by wahid18-maqsLoading…
Fix the Attribute Error in gradient_accumulation_and_microbatching.ipynb example notebook
#1740 opened Jul 28, 2026 by rajasekharporeddyCollaboratorLoading…
fix: warn when Hutchinson estimator uses fixed PRNGKey(0)
#1738 opened Jul 27, 2026 by wahid18-maqsLoading…
docs: fix COCOB docstring reference and author name
#1737 opened Jul 27, 2026 by wahid18-maqsLoading…
Support scalar parameters in scale_by_sm3
#1736 opened Jul 26, 2026 by TheSaiEarantiLoading…
Compute exact matrix powers in matrix_inverse_pth_root for any p
#1735 opened Jul 26, 2026 by TheSaiEarantiLoading…
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-serviceBotLoading…
1
docs(contrib): fix paper reference link for ADOPT optimizer
#1731 opened Jul 22, 2026 by himotoLoading…
1
Treat integer hyperparameters as static in inject_hyperparams (fixes #412)
#1730 opened Jul 21, 2026 by ArneshBanerjeeLoading…
1
1
Add nbstripout pre-commit hook, keeping notebook outputs (#1372)
#1729 opened Jul 20, 2026 by shubhangithubLoading…
1
1
Use jax.Array instead of jax.typing.ArrayLike across the codebase.
#1728 opened Jul 14, 2026 by copybara-serviceBotLoading…
Add optional batch shape argument to tree_random_like.
#1727 opened Jul 13, 2026 by carlosgmartinContributorLoading…
Informational, do not merge
#1726 opened Jul 13, 2026 by copybara-serviceBotLoading…
1
Fix scale_by_rprop returning a one-step-stale update
#1725 opened Jul 12, 2026 by winklemadLoading…
Fix projection_box crash on scalar bounds with multi-leaf pytrees
#1724 opened Jul 12, 2026 by winklemadLoading…
add missing docstring to get_spectral_radius_upper_bound in linear_algebra.py
#1721 opened Jul 4, 2026 by RavSinghChandanLoading…
9
add missing docstrings to tile_second_to_last_dim, multi_normal and parse_version
#1720 opened Jul 4, 2026 by RavSinghChandanLoading…
4
contrib: add SCION optimizer (norm-constrained LMO generalisation of Muon)
#1718 opened Jun 27, 2026 by irhylContributorLoading…
1
Fix ACProp second moment to use the new first moment
#1716 opened Jun 27, 2026 by codewithfourtixLoading…
1
Fix incorrect (1 - b2) coefficient in scale_by_adan
#1715 opened Jun 26, 2026 by codewithfourtixLoading…
5
ProTip! Filter pull requests by the default branch with base:main.