-
Notifications
You must be signed in to change notification settings - Fork 339
Pull requests: google-deepmind/optax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
contrib: add SCION optimizer (norm-constrained LMO generalisation of Muon)
#1718
opened Jun 27, 2026 by
irhyl
Contributor
Loading…
Fix ACProp second moment to use the new first moment
#1716
opened Jun 27, 2026 by
codewithfourtix
Loading…
Fix incorrect (1 - b2) coefficient in scale_by_adan
#1715
opened Jun 26, 2026 by
codewithfourtix
Loading…
feat: Add SPSA optimization method (Issue #357)
#1712
opened Jun 23, 2026 by
Paramveersingh-S
Loading…
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…
fix: memoize optimizer factory functions to avoid JIT recompilation (#353)
#1693
opened Jun 9, 2026 by
Muneerali199
Contributor
Loading…
contrib: add SOAP optimizer (arXiv:2409.11321)
#1692
opened Jun 9, 2026 by
irhyl
Contributor
Loading…
feat(contrib): add cautious optimizer wrapper
#1690
opened Jun 5, 2026 by
Sumu004
Loading…
4 tasks done
feat(contrib): add MARS optimizer (variance-reduction AdamW)
#1689
opened Jun 5, 2026 by
Sumu004
Loading…
fix(transform): cast updates to grad dtype for mixed-precision support
#1687
opened May 30, 2026 by
nileshpatil6
Loading…
fix(inject_hyperparams): treat plain Python ints as static to fix adafactor jit (#412)
#1686
opened May 30, 2026 by
nileshpatil6
Loading…
Fix pytype and runtime errors in Optax causing GitHub CI failures.
#1681
opened May 26, 2026 by
copybara-service
Bot
Loading…
Resolve various attribute access and overload warnings in losses, transforms, and contributed optimizers.
#1678
opened May 23, 2026 by
copybara-service
Bot
Loading…
feat(perturbations): rename sigma→scale, default num_samples=1 (#1342)
#1662
opened Apr 19, 2026 by
yurekami
Loading…
2 of 3 tasks
Issue #1658 - Add SSBroyden and SSBFGS optimizers to optax.contrib
#1659
opened Apr 10, 2026 by
armbrusl
Loading…
Add NorMuon optimizer (row-wise adaptive normalization for Muon)
#1651
opened Mar 30, 2026 by
bzantium
Contributor
Loading…
4 tasks done
Remove deprecated second_order functions for 0.2.9
#1646
opened Mar 23, 2026 by
dylantirandaz
Loading…
2 tasks
Extend muon to support mask_callable weight dimension numbers
#1643
opened Mar 23, 2026 by
dylantirandaz
Loading…
3 tasks
Disable pmap-shard_map merge if it's available.
#1636
opened Mar 20, 2026 by
copybara-service
Bot
Loading…
Previous Next
ProTip!
no:milestone will show everything without a milestone.