Skip to content

Pull requests: google-deepmind/optax

New pull request New
37 Open 1,110 Closed
37 Open 1,110 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

Add madgrad optimizer
#1581 opened Jan 25, 2026 by starryendymion Loading…
1
5
Upgrade GitHub Actions to latest versions
#1578 opened Jan 25, 2026 by salmanmkc Loading…
clipping: support default unitwise_norm for 5D params
#1576 opened Jan 24, 2026 by staticpayload Loading…
Add example demonstrating the microbatching api and comparing it against optax.MultiSteps.
#1573 opened Jan 23, 2026 by copybara-service bot Loading…
Upgrade GitHub Actions for Node 24 compatibility
#1572 opened Jan 22, 2026 by salmanmkc Loading…
Remove TF dependency in Lookahead Optimizer on MNIST example
#1568 opened Jan 22, 2026 by rajasekharporeddy Loading…
1
implements Mars optimizer [contrib]
#1562 opened Jan 13, 2026 by KumarADITHYA123 Loading…
1
2
Raise error on unused extra kwargs in backtracking linesearch
#1559 opened Jan 12, 2026 by TanmayThakur2209 Loading…
7
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…
chore: Remove deprecated schedules/inject.py stub file
#1548 opened Jan 6, 2026 by harshaljanjani Loading…
4
Include py.typed marker for PEP 561 compliance (#1388)
#1546 opened Jan 4, 2026 by Ashutosh0x Loading…
1
Add sparse simplex, transport, and Birkhoff projections from JAXopt
#1545 opened Jan 3, 2026 by aymuos15 Loading…
7
Add example for correct usage of Optax lookahead optimizer
#1542 opened Dec 28, 2025 by Solventerritory Loading…
5
Remove TensorFlow dependency and migrate mlp_mnist to Flax NNX
#1536 opened Dec 18, 2025 by selamw1 Loading…
3
Update self-supervised losses implementation and tests
#1531 opened Dec 11, 2025 by surajyadav-research Loading…
30
docs: Add missing alt text to gallery images for accessibility.
#1523 opened Dec 5, 2025 by Sandeep3084 Loading…
9
Update make_perturbed_fun
#1520 opened Dec 3, 2025 by haroon0x Loading…
1
1
Feat: Add pre-commit hook that strips notebooks of metadata
#1519 opened Dec 3, 2025 by haroon0x Loading…
1
11
Fix: #1509 Merge duplicate DoG implementations and add layer-wise support
#1518 opened Dec 2, 2025 by abheeyeee Loading…
29
Remove TensorFlow dependency and migrate cifar10_resnet to Flax NNX
#1508 opened Nov 20, 2025 by selamw1 Loading…
4
Adds utilities for reshaping Optax states to separate parameter-like fields.
#1495 opened Nov 14, 2025 by copybara-service bot Loading…
Add aggregators in optax.
#1432 opened Oct 18, 2025 by copybara-service bot Loading…
Minor improvements and cleanup using tree utility functions.
#1365 opened Jun 27, 2025 by carlosgmartin Loading…
2
ProTip! Type g i on any issue or pull request to go back to the issue listing page.