Skip to content

Pull requests: google-deepmind/optax

New pull request New
62 Open 1,219 Closed
62 Open 1,219 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

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…
1
add missing docstrings to tile_second_to_last_dim, multi_normal and parse_version
#1720 opened Jul 4, 2026 by RavSinghChandan Loading…
1
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…
Migrate the Adversarial training example to the latest flax.nnx module.
#1710 opened Jun 22, 2026 by rajasekharporeddy Collaborator Loading…
Migrate to the latest flax.nnx module from the legacy flax.linen module in Character-level Transformer on Tiny Shakespeare example.
#1709 opened Jun 19, 2026 by rajasekharporeddy Collaborator Loading…
Add experimental sharding module with with_custom_sharding for
#1704 opened Jun 14, 2026 by copybara-service Bot Loading…
Add affine set and box section projections from JAXopt
#1701 opened Jun 13, 2026 by discobot Loading…
1
1
Add self-supervised losses
#1695 opened Jun 9, 2026 by ahmedtaha100 Loading…
1
1
fix: memoize optimizer factory functions to avoid JIT recompilation (#353)
#1693 opened Jun 9, 2026 by Muneerali199 Contributor Loading…
1
2
contrib: add SOAP optimizer (arXiv:2409.11321)
#1692 opened Jun 9, 2026 by irhyl Contributor Loading…
2
feat(contrib): add cautious optimizer wrapper
#1690 opened Jun 5, 2026 by Sumu004 Loading…
4 tasks done
1
feat(contrib): add MARS optimizer (variance-reduction AdamW)
#1689 opened Jun 5, 2026 by Sumu004 Loading…
1
3
fix(transform): cast updates to grad dtype for mixed-precision support
#1687 opened May 30, 2026 by nileshpatil6 Loading…
1
2
fix(inject_hyperparams): treat plain Python ints as static to fix adafactor jit (#412)
#1686 opened May 30, 2026 by nileshpatil6 Loading…
1
Fix pytype and runtime errors in Optax causing GitHub CI failures.
#1681 opened May 26, 2026 by copybara-service Bot Loading…
ProTip! Mix and match filters to narrow down what you’re looking for.