Skip to content

Pull requests: google-deepmind/optax

New pull request New
47 Open 1,151 Closed
47 Open 1,151 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

Disable pmap-shard_map merge if it's available.
#1636 opened Mar 20, 2026 by copybara-service bot Loading…
Standardize weight_decay_mask naming and add weight decay
#1635 opened Mar 20, 2026 by rdyro Loading…
feat: add optional weight_decay to adabelief optimizer
#1634 opened Mar 19, 2026 by yasumorishima Loading…
3 tasks done
1
1
internal change
#1631 opened Mar 17, 2026 by copybara-service bot Loading…
Adding pre-commit action to build markdown files derived from example jupyter notebooks.
#1627 opened Mar 12, 2026 by rdyro Loading…
experimental notebook markdown preview generation via actions
#1626 opened Mar 12, 2026 by rdyro Loading…
Update mlp_mnist notebook for removing tf dependency and NNX migration
#1625 opened Mar 12, 2026 by selamw1 Loading…
@emilyfertig
1
Enable performance gain by batching small update tensors.
#1618 opened Mar 5, 2026 by copybara-service bot Loading…
1
Implement SOAP optimizer (#1580)
#1617 opened Mar 5, 2026 by jinukuntlaakhilakumargoud-web Loading…
1
Remove TensorFlow dependency in Reduce on Plateau Learning Rate Scheduler example
#1616 opened Mar 5, 2026 by rajasekharporeddy Loading…
Rename log_prob to unnormalized_log_prob in perturbations (#1214)
#1614 opened Mar 5, 2026 by jinukuntlaakhilakumargoud-web Loading…
1
3
Add PolarExpress variant to Muon optimizer
#1613 opened Mar 4, 2026 by MarcMachaczek Loading…
1
5
Add NorMuon optimizer
#1612 opened Feb 28, 2026 by tinker495 Loading…
3 tasks done
1
Add cautious weight decay transformation
#1611 opened Feb 28, 2026 by tinker495 Loading…
3 tasks done
@vroulet
2
Remove TensorFlow dependency in the differentially private CNN on MNIST example
#1610 opened Feb 26, 2026 by rajasekharporeddy Loading…
1
fix: make scale_by_distance_over_gradients a deprecated wrapper aroun…
#1606 opened Feb 22, 2026 by Mahajan-Sachin Loading…
1
10
Adahessian optimizer and refactor Hutchinson estimator as shared utility to be used by Adahessian, Sophia optimizer
#1604 opened Feb 19, 2026 by viralvgupta Loading…
19
fix: correct typo 'seperately' to 'separately'
#1599 opened Feb 9, 2026 by thecaptain789 Loading…
2
Add asymmetric loss
#1591 opened Jan 31, 2026 by DavidRConnell Loading…
1
Migrate from flax.linen to flax.nnx in Lookahead Optimizer in MNIST example.
#1589 opened Jan 29, 2026 by rajasekharporeddy Loading…
implements Mars optimizer [contrib]
#1562 opened Jan 13, 2026 by KumarADITHYA123 Loading…
1
2
Add regression test for scale_by_rms zero-gradient stability
#1553 opened Jan 7, 2026 by TanmayThakur2209 Loading…
2
Add semi-supervised losses and test
#1551 opened Jan 6, 2026 by surajyadav-research Loading…
9
[graphcast] Prepare for jax_pmap_shmap_merge=True.
#1549 opened Jan 6, 2026 by copybara-service bot Loading…
Include py.typed marker for PEP 561 compliance (#1388)
#1546 opened Jan 4, 2026 by Ashutosh0x Loading…
1
ProTip! Follow long discussions with comments:>50.