Skip to content

Pull requests: google-deepmind/optax

New pull request New
39 Open 1,061 Closed
39 Open 1,061 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 auction algorithm implementation and tests
#1521 opened Dec 3, 2025 by surajyadav-research Loading…
1
3
Update make_perturbed_fun
#1520 opened Dec 3, 2025 by haroon0x Loading…
1
Feat: Add pre-commit hook that strips notebooks of metadata
#1519 opened Dec 3, 2025 by haroon0x Loading…
1
2
Fix: #1509 Merge duplicate DoG implementations and add layer-wise support
#1518 opened Dec 2, 2025 by abheeyeee Loading…
12
[optax] Remove jax_pmap_shmap_merge=False flag in contrib/_complex_valued. #jax-fixit
#1517 opened Dec 2, 2025 by copybara-service bot Loading…
1
Clarify difference between kl_divergence and convex_kl_divergence
#1514 opened Nov 29, 2025 by zer-art Loading…
1 task done
1
16
Add micro_value_and_grad to microbatching.py
#1510 opened Nov 24, 2025 by copybara-service bot Loading…
Remove TensorFlow dependency and migrate cifar10_resnet to Flax NNX
#1508 opened Nov 20, 2025 by selamw1 Loading…
2
Fix conjugation in Newton-Schulz iterator and update tests for comple…
#1506 opened Nov 20, 2025 by eirikfagerbakke Loading…
1
2
Fix multiply_by_parameter_scale type in adafactor optimizer.
#1504 opened Nov 18, 2025 by copybara-service bot Loading…
Optimize tree_sum compile time using tree_reduce_associative
#1503 opened Nov 18, 2025 by Aaryan-549 Loading…
4
Added Refined Lion Optimizer and Tests
#1497 opened Nov 16, 2025 by raghulchandramouli Loading…
18
Adds utilities for reshaping Optax states to separate parameter-like fields.
#1495 opened Nov 14, 2025 by copybara-service bot Loading…
Add axis argument to smooth_labels function.
#1492 opened Nov 12, 2025 by carlosgmartin Loading…
Add an internal definition of ArrayTree and use it instead of chex.ArrayTree.
#1484 opened Nov 12, 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
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
ProTip! Follow long discussions with comments:>50.