Skip to content

Pull requests: google-deepmind/optax

New pull request New
32 Open 1,046 Closed
32 Open 1,046 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 axis argument to smooth_labels function.
#1492 opened Nov 12, 2025 by carlosgmartin Loading…
Enforce keyword-only arguments for optional parameters in Optax losses.
#1490 opened Nov 12, 2025 by copybara-service bot Loading…
Define PRNGKey and Shape in optax._src.base and update usages.
#1489 opened Nov 12, 2025 by copybara-service bot Loading…
Add an internal definition of ArrayTree and use it instead of chex.ArrayTree.
#1484 opened Nov 12, 2025 by copybara-service bot Loading…
Replace chex.Numeric with jax.typing.ArrayLike. Replace chex.Scalar with float/int as appropriate.
#1479 opened Nov 11, 2025 by copybara-service bot Loading…
Remove use of cast_tree in favor of optax.tree.cast
#1477 opened Nov 11, 2025 by copybara-service bot Loading…
Remove not exposed file.
#1474 opened Nov 11, 2025 by copybara-service bot Loading…
Add an internal definition of ArrayTree and use it instead of chex.ArrayTree.
#1470 opened Nov 10, 2025 by copybara-service bot Loading…
Fix typo in alias.py documentation
#1464 opened Nov 7, 2025 by partev Loading…
1
add Dice + CE loss and weighting to current Dice implementation
#1458 opened Nov 2, 2025 by aymuos15 Loading…
1
2
fix: default hyperparams to float32 unless specified
#1445 opened Oct 25, 2025 by akramcodez Loading…
1
2
Fix _projections_test.py: Remove prints and use optax.tree.allclose.
#1439 opened Oct 24, 2025 by carlosgmartin Loading…
1
feat - Implemented Lion Optimizer and Tests scripts
#1438 opened Oct 24, 2025 by raghulchandramouli Loading…
2
Add aggregators in optax.
#1432 opened Oct 18, 2025 by copybara-service bot Loading…
Add a bias_correction_v flag to scale_by_amsgrad to align with the original AMSGrad paper and Pytorch/tensorflow impl
#1423 opened Sep 26, 2025 by vvsvictor Loading…
1
7
Implementing ANO(All new Optimizer) in Optax.
#1416 opened Sep 4, 2025 by tinker495 Loading…
4
Tests to warmup_cosine_decay_schedule edge cases
#1413 opened Sep 1, 2025 by edawite Loading…
15
Fix tree_min and tree_max to handle zero-size array leafs.
#1374 opened Jul 6, 2025 by carlosgmartin Loading…
3
Minor improvements and cleanup using tree utility functions.
#1365 opened Jun 27, 2025 by carlosgmartin Loading…
2
Add tree_bits function.
#1341 opened Jun 18, 2025 by carlosgmartin Loading…
20
Silence some pytype errors related to a JAX build refactor
#1338 opened Jun 16, 2025 by copybara-service bot Loading…
Avoid unnecessary directory change in test.sh.
#1324 opened Jun 9, 2025 by carlosgmartin Loading…
Add weight_decay and mask arguments to adabelief optimizer.
#1291 opened May 1, 2025 by carlosgmartin Loading…
1
6
Add support for loss functions with auxiliary data to linesearch
#1177 opened Jan 16, 2025 by ro0mquy • Draft
4 tasks
10
Add linear program solver based on the restarted Halpern primal-dual hybrid gradient (rHPDHG) algorithm.
#1154 opened Dec 9, 2024 by carlosgmartin Loading…
1
ProTip! Updated in the last three days: updated:>2025-11-09.