Optax 0.2.4

What's Changed

• Beginning of 0.2.4 development by @copybara-service in #1003
• Fix gallery display, add images by @vroulet in #1005
• Fix docs for dog by @jungtaekkim in #1008
• docs: remove multi_normal from utilities.rst by @Abhinavcode13 in #1009
• Fix docs for dowg by @jungtaekkim in #1013
• feat: add a mathematical description of AdaGrad optimizer by @Abhinavcode13 in #1011
• fix: refactor AdaGrad optimizer recent changes by @Abhinavcode13 in #1016
• Make masking compatible with callable pytrees a la Equinox by @copybara-service in #1015
• Enable zero lrate for schedule-free optimization by @copybara-service in #1018
• keep a local .pylintrc file by @fabianp in #1024
• Add bias_correction and eps_in_sqrt options to rmsprop and associated transforms by @copybara-service in #1019
• Replace adam(b1=0) by rmsprop for schedule_free by @copybara-service in #1025
• Update init value for zakharov problem from 1e4 to 1e3 by @copybara-service in #1027
• Fix typo by @gil2rok in #1030
• Updated docs cosine_decay_schedule by @bhargavyagnik in #1032
• feat: add mathematical notation docs of SM3 optimizer by @Abhinavcode13 in #1012
• DOC: misc improvements in docstring of softmax_cross_entropy* by @copybara-service in #1033
• add doctest to constant_schedule by @fabianp in #1034
• Add axis and where arguments to loss functions. by @carlosgmartin in #912
• Fix doctest error in make_fenchel_young_loss by @copybara-service in #1035
• add doctest for polynomial_schedule by @fabianp in #1037
• add missing schedule_free_* methods by @fabianp in #1043
• fix error in softmax_cross_entropy formula by @fabianp in #1041
• Fix typo in formula of cosine_decay_schedule by @fabianp in #1044
• schedule_free: fix broadcasting of scalar arrays to 1d arrays by @n-gao in #1042
• Update polynomial_schedule doctest per @vroulet's feedback by @fabianp in #1045
• Fix linting schedule_free_test by @copybara-service in #1048
• more robust tests by @copybara-service in #1050
• Generalizes safe_int32_increment to safe_increment by @copybara-service in #1054
• Add dtype option to tree_random_like by @copybara-service in #1056
• Add double precision tests for safe_increment and fix warnings on float64_test.py by @copybara-service in #1055
• Add optax.tree_utils.tree_random_split. by @carlosgmartin in #1063
• Fix test.sh, which uses set -u, so that it works when JAX_VERSION is unset. by @carlosgmartin in #1070
• Migrate from jax.tree_util legacy APIs to new jax.tree API. by @carlosgmartin in #1066
• Ensure optimizers return updates of same dtype as params. by @copybara-service in #1060
• Fix test.sh to not modify .pylintrc. by @carlosgmartin in #1071
• Replace deprecated typing.Hashable with collections.abc.Hashable. by @carlosgmartin in #1068
• Relax absolute tolerance for failing tests involving chex.assert_trees_all_close. by @carlosgmartin in #1069
• Fix doctests by @copybara-service in #1073
• Tidy up test.sh and make it clean up properly. by @carlosgmartin in #1074
• Add missing initializer argument of 0 to tree_reduce in tree_vdot and tree_sum. by @carlosgmartin in #1065
• bump chex version for #1076 by @fabianp in #1078
• correct RST references by @fabianp in #1079
• deprecate methods in optax/monte_carlo. by @copybara-service in #1076
• schedule-free optimier: ensure it's possible to donate both the state and the params by @enolan in #1059
• add link to examples from docstring by @fabianp in #1085
• Adopt US spelling for documentation and fix typos by @miguelcsx in #1087
• Update docs: Note RMSprop usage instead of Adam for memory savings in… by @nasyxx in #1086
• adding a perturbations module. Can take pytrees as inputs by @copybara-service in #827
• Fix initial step of scale_by_optimistic_gradient. by @carlosgmartin in #1088
• Add Hungarian algorithm for the linear assignment problem. by @carlosgmartin in #1083
• Allow safe_increment to handle unsigned integers. by @carlosgmartin in #1092
• Fix formatting issues with gallery entry for linear assignment problem. by @carlosgmartin in #1095
• use sphinx references instead of hardcoded links. by @fabianp in #1096
• Remove dtype safeguards by @copybara-service in #1099
• cosmetic improvements perturbations module by @fabianp in #1097
• update jax.tree.map to comply with jax 0.4.34 by @a1302z in #1094
• Add Adan optimizer. by @carlosgmartin in #1090
• Fix typo in projection_simplex docstring. by @copybara-service in #1105
• add config for link checking, and mark 429 (too many requests) as fine by @fabianp in #1103
• Fix docstring for hungarian_algorithm. by @carlosgmartin in #1102
• Add optax.optimistic_adam. by @carlosgmartin in #1089
• Add projection_l1_sphere and projection_l1_ball. by @copybara-service in #1106
• Add projection_l2_sphere and projection_l2_ball. by @copybara-service in #1114
• Add tree_max. by @copybara-service in #1115
• Add projection_linf_ball. by @copybara-service in #1117
• remove test that leaked jax tracers by @copybara-service in #1123
• Add a mathematical description for Lion by @aman2304 in #1121
• Fix the sign for the update in the math equation for nadam in the docs. by @carlosgmartin in #1128
• ntxent fix by @GrantMcConachie in #946
• Add Nesterov momentum to AdaBelief optimizer. by @carlosgmartin in #1127
• fix: Coherent dtypes of updates with and without MultiSteps by @hlzl in #1122
• Fix AdaBelief implementation. by @carlosgmartin in #1130
• Revisiting linesearches and LBFGS. by @copybara-service in #1133

New Contributors

• @vroulet made their first contribution in #1005
• @jungtaekkim made their first contribution in #1008
• @Abhinavcode13 made their first contribution in #1009
• @gil2rok made their first contribution in #1030
• @bhargavyagnik made their first contribution in #1032
• @n-gao made their first contribution in #1042
• @enolan made their first contribution in #1059
• @miguelcsx made their first contribution in #1087
• @a1302z made their first contribution in #1094
• @aman2304 made their first contribution in #1121
• @hlzl made their first contribution in #1122

Full Changelog: https://githu...

Optax 0.2.3

What's Changed

• Fix the KeyboardInterrupt exception from #860 by removing the timeout by @copybara-service in #886
• Beginning of 0.2.3 development by @copybara-service in #893
• Add a mathematical description of AdamW by @gbruno16 in #894
• Suppress not-callable pylint error for now since is being flagged erroneously all over the place. by @copybara-service in #908
• Fix doc link by @yixiaoer in #903
• Fixed pseudocode for Nesterov in description of SGD. by @satyenkale in #901
• Fix softmax_cross_entropy to handle -inf logits correctly when corresponding label is 0. by @carlosgmartin in #898
• Upstream sparsemax jaxopt loss to optax. by @copybara-service in #899
• Reorganize tree_utils. by @copybara-service in #914
• Revert of #898. by @copybara-service in #915
• Fix jax.tree_map deprecation warnings. by @copybara-service in #917
• Correct handling of -inf in softmax_cross_entropy. Fix #898. by @copybara-service in #916
• Added mathematical documentation to AdaMax by @hmludwig in #918
• Fix pip install command for doc dependencies. by @mblondel in #922
• Start documentation for projections. by @mblondel in #921
• Add projection_simplex. by @copybara-service in #919
• Move gradient transformations to optax.transforms sub-package - 1/N by @copybara-service in #923
• Added a NTXent loss by @GrantMcConachie in #897
• fix(docs): broken link in README by @jeertmans in #940
• Add a deprecation module to warn or raise errors for deprecations (following jax semantics). by @copybara-service in #931
• chore(ci): add markdown-link-check action by @jeertmans in #939
• Implementation of MoMo algorithm by @fabian-sp in #721
• Weight decay for COCOB by @albcab in #945
• Add a nesterov flag to radam optimizer. by @carlosgmartin in #949
• Formatting in momo docstring + doctest by @fabianp in #950
• docstring formatting by @fabianp in #952
• Port schedule_free optimizer to optax. Original pytorch repo: https://github.com/facebookresearch/schedule_free by @copybara-service in #911
• Fix RST formatting issues. by @fabianp in #953
• remove duplicated BATCH_SIZE argument by @fabianp in #956
• Replace deprecated jax.tree_* functions with jax.tree.* by @copybara-service in #963
• remove residues from previous builds before running tests by @fabianp in #967
• Fix docs errors by @copybara-service in #941
• Removing sophia optimizer by @copybara-service in #973
• move clipping transforms to optax.transforms. by @copybara-service in #926
• Expose components in sub-package by @copybara-service in #978
• Add multiclass_sparsemax_loss. by @copybara-service in #971
• Remove useless inner jit by @copybara-service in #957
• Fix memory leak in radam optimizer by @lukekulik in #974
• Add end_scale argument by @stefanocortinovis in #975
• Fix error with x64 loss by @stefanocortinovis in #976
• LBFGS solver part 1: chainable preconditioner. by @copybara-service in #980
• Fix docs errors (following warnings displayed in doc logs of github actions) by @copybara-service in #984
• [JAX] Update users of jax.tree.map() to be more careful about how they handle Nones. by @copybara-service in #983
• LBFGS solver part 2: implementing linesearch ensuring sufficient decrease and small curvature by @copybara-service in #981
• CI: add test against oldest supported JAX version by @jakevdp in #987
• Internal change by @copybara-service in #988
• Ignore some linesearch tests on gpu/tpu by @copybara-service in #986
• LBFGS part 3: combine lbfgs and zoom linesearch by @copybara-service in #989
• Add arxiv reference to schedule_free optimizer. by @copybara-service in #997
• LBFGS part 4: notebook illustrating how to use lbfgs with linesearch as a solver. by @copybara-service in #991
• Add common schedule_free wrappers. by @copybara-service in #998
• Add schedule_free check for b1 != 0. by @copybara-service in #999
• feat: add normalize_by_update_norm by @SauravMaheshkar in #958
• Saurav maheshkar saurav/scale by grad norm by @fabianp in #1000
• Fix doctest normalize_by_update_norm by @copybara-service in #1002
• Release v0.2.3 by @copybara-service in #1001

New Contributors

• @gbruno16 made their first contribution in #894
• @satyenkale made their first contribution in #901
• @GrantMcConachie made their first contribution in #897
• @fabian-sp made their first contribution in #721
• @lukekulik made their first contribution in #974
• @jakevdp made their first contribution in #987

Full Changelog: v0.2.2...v0.2.3

Optax 0.2.2

What's Changed

• Added mathematical description to Noisy SGD by @hmludwig in #857
• Use sphinx-contributors for an automated contributors list. by @fabianp in #841
• Implementation of the Polyak SGD solver by @copybara-service in #718
• Document the extra args of the update function in docstring by @copybara-service in #864
• Utility to set value in a pytree (and so in state) by @copybara-service in #865
• Added mathematical description to AdaBelief docstring by @hmludwig in #869
• FIX RST formatting in inject hyperparams by @fabianp in #867
• Warn that in future arguments after the initial (prediction, ground_truth) positional arguments will become keyword-only in optax losses. by @copybara-service in #863
• Upstream missing jaxopt losses to optax - Part 2/N by @copybara-service in #872
• Fix error reduce_on_plateau.ipynb:20002: WARNING: No source code lexer found for notebook cell by @copybara-service in #875
• docstring cosmetic improvements by @fabianp in #879
• Extend capabilities of tree_get, tree_set. by @copybara-service in #878
• [DOC] Add to the gallery an example on a small language model by @copybara-service in #866
• Update reduce_on_plateau to handle training average loss. by @copybara-service in #883
• Fix notebook reduce_on_plateau by @copybara-service in #887
• ENH: extend power_iteration to accept a matrix in implicit form by @copybara-service in #858
• Document changes in power_iteration by @copybara-service in #889
• Release of version 0.2.2 by @copybara-service in #892

Full Changelog: v0.2.1...v0.2.2

v0.2.1

What's Changed

• Begin of 0.2.1 development by @copybara-service in #845
• Add mathematical formula in docstring of linear_schedule by @copybara-service in #839
• Reference reduce on plateau example from docstring by @copybara-service in #851
• Fix sign discrepancy in tolerance in scale_by_backtracking_linesearch by @copybara-service in #853
• 0.2.1 release by @copybara-service in #855

Full Changelog: v0.2.0...v0.2.1

v0.2.0

What's Changed

• Begin of 0.2.0 development by @copybara-service in #766
• Remove unnecessary %pip install statements from the notebooks by @JayJayleee in #762
• Updates warmup_cosine_decay_schedule to allow 0 as peak_value. Currently it errors out as divide by 0. by @copybara-service in #770
• fixed math flex direction by @copybara-service in #775
• Docstrings for all solvers with Objective Function values by @copybara-service in #771
• Misc improvements in adversarial_training.ipynb by @copybara-service in #776
• Solver API (wip) by @copybara-service in #777
• Detail that learning_rate can be scalar or a schedule. by @copybara-service in #778
• workaround for pylint issues by @fabianp in #783
• correct type annotation for params by @yixiaoer in #785
• Ensure no execution of slow notebooks. by @copybara-service in #788
• [DOC] merge development and design_docs into a single file by @fabianp in #789
• Add development information in README. by @mblondel in #790
• Solve #767, some examples not being displayed on the webpage by @fabianp in #791
• Add missing examples to contrib examples section. by @fabianp in #794
• Make EmptyState a valid ArrayTree by @copybara-service in #752
• Improve error message for cosine_decay_schedule. by @copybara-service in #797
• move DPSGD docs to contrib/ section by @WangHanSolo in #782
• Adjusted 'test.sh' for universal macOS and Linux Compatibility by @hmludwig in #804
• Benchmark recommendations. by @copybara-service in #807
• Removed 'examples/quick_start.ipynb' and adjusted references. by @hmludwig in #811
• [DOC] streamline sidebar of gallery by @fabianp in #810
• [FIX] bad indentation in the docstring of nadam by @fabianp in #812
• Backtracking linesearch. by @copybara-service in #795
• Fixing backtracking linesearch support on gpus by @copybara-service in #814
• allow wrappers.flatten to work with 0-dim arrays by @copybara-service in #816
• Add Example, Args and Returns section to optax.linear_schedule by @copybara-service in #818
• FIX doctest format in TransformUpdateExtraArgsFn by @copybara-service in #820
• Refactoring and link fixes in development.md by @fabianp in #822
• Clarify formula for cosine learning rate. by @copybara-service in #823
• Detail SGD with momentum equations. by @copybara-service in #830
• Fixed the rendering of equations in jupyter notebooks by @hmludwig in #836
• Fixed the Sharpness aware minimization docstring by @hmludwig in #834
• FIX "gallery.rst:5: WARNING: Title underline too short." by @copybara-service in #838
• Moved linear algebra operations into utilities by @hmludwig in #837
• adding tree utils for generating trees with random values and averaging across trees. by @copybara-service in #825
• v0.2.0 release by @copybara-service in #844

New Contributors

• @JayJayleee made their first contribution in #762
• @yixiaoer made their first contribution in #785
• @mblondel made their first contribution in #790
• @WangHanSolo made their first contribution in #782
• @hmludwig made their first contribution in #804

Full Changelog: v0.1.9...v0.2.0

v0.1.9

What's Changed

• update URL github.com/deepmind -> github.com/google-deepmind and branch to main by @copybara-service in #710
• Show CI results only for the main branch by @fabianp in #716
• typo nesterov -> Nesterov by @fabianp in #722
• Add atol option to contrib.reduce_on_plateau() by @stefanocortinovis in #698
• add docs for tree_utils module by @amosyou in #724
• Simplifications on the doc by @copybara-service in #727
• add nadam and nadamw optimizers by @copybara-service in #723
• Add versionadded and seealso metadata to (n)adam(w) solvers by @copybara-service in #729
• Enable doctests in sphinx and fix failing doctests by @fabianp in #733
• Add missing members to the doc by @copybara-service in #734
• FIX sphinx warnings "this document is not included in any toctree" by @fabianp in #736
• feat(ci): drop setup.py from publishing CI by @SauravMaheshkar in #737
• Minor tweak in pypi-publish.yml by @fabianp in #739
• [TEST] Install virtual environment in current directory instead of /tmp by @fabianp in #746
• migrate to pyproject by @copybara-service in #747
• Deprecate optax.inject_stateful_hyperparameters and replace it with optax.inject_hyperparameters by @copybara-service in #730
• Clarify inclusion criteria into optax and optax.contrib by @copybara-service in #742
• fix the default learning rate in prodigy by @konstmish in #740
• update and merge quickstart notebooks by @amosyou in #726
• Remove redundant examples from README by @mmhamdy in #754
• Instructions to build the doc and option to build the docs without executing the notebooks. by @copybara-service in #759
• Remove license statement from notebooks by @renuka010 in #764

New Contributors

• @amosyou made their first contribution in #724
• @konstmish made their first contribution in #740
• @renuka010 made their first contribution in #764

Full Changelog: v0.1.8...v0.1.9

Optax 0.1.8

What's Changed

• Remove legacy symbols. by @copybara-service in #397
• Move second order utilities to dedicated sub-package. by @copybara-service in #588
• Move stochastic gradient estimation functions to separate monte_carlo subpackage. by @copybara-service in #589
• Move losses to dedicated sub-package. by @copybara-service in #590
• Note TF's adam implementation difference by @vwxyzjn in #572
• Note RMSProp PyTorch vs Optax implementation difference by @vwxyzjn in #595
• Move complex_valued from experimental/ to contrib/. by @copybara-service in #602
• Make radam docstring point to radam paper. by @carlosgmartin in #510
• COntinuous COin Betting algorithm by @albcab in #536
• Update second_order package internal structure. by @copybara-service in #625
• Update github URL + nicer console highlighting by @copybara-service in #628
• Add import stub. by @copybara-service in #633
• Start optax.projections subpackage and add projection_non_negative. by @copybara-service in #632
• Also apply masking to compatible extra_args. by @copybara-service in #620
• Add tree_vdot to tree_utils. by @copybara-service in #634
• Add: merge multistep and apply_every logic by @celiolarcher in #596
• Remove redundant requirements folder by @copybara-service in #643
• Run pytype on Python 3.11 by @fabianp in #657
• Move MLP MNIST example to Colab notebook by @copybara-service in #658
• Move DP-SGD example to Colab by @copybara-service in #660
• Add tensorflow as dependency for the examples by @copybara-service in #661
• Add tensorflow-datasets to doc requirements by @copybara-service in #662
• D-Adaptation and Prodigy contrib implementations by @adefazio in #651
• Fix typing issue in prodigy.py for python 3.9 by @copybara-service in #671
• Fix CIFAR10 Flax&Optax example name in the gallery by @copybara-service in #672
• Move flax example to Colab by @copybara-service in #670
• update citation url since http://github.com/deepmind is now empty by @copybara-service in #667
• Move Haiku example to Colab notebook by @copybara-service in #673
• run pytype also on the contrib directory by @copybara-service in #675
• Add reduce_on_plateau LR scheduler to contrib directory. by @vz415 in #629
• Add to .gitignore temporary directories of building the docs. by @fabianp in #680
• Add the SAM example to the example gallery by @fabianp in #681
• Add reduce_on_plateau example (#679) by @copybara-service in #683
• Wrong image path on redue on plateau example by @copybara-service in #686
• Remove execution timeout (#687) by @copybara-service in #689
• Documentation Typo Fixes for apply_if_finite by @mharradon in #696
• FIX typos in public readme by @copybara-service in #691
• Add adversarial training example by @mmhamdy in #682
• Divide API documentation by @MaanasArora in #688
• Avoid scalar conversion of non-scalar arrays by @copybara-service in #693
• Improving docstrings for schedules. by @copybara-service in #692
• Fix TypeError in contrib.reduce_on_plateau() when x64 is enabled by @stefanocortinovis in #697
• Fix some doc links in README.md and other files #665 by @Bigpet in #706
• replace branch master by main by @copybara-service in #709
• New optax release by @copybara-service in #701

New Contributors

• @vwxyzjn made their first contribution in #572
• @carlosgmartin made their first contribution in #510
• @albcab made their first contribution in #536
• @celiolarcher made their first contribution in #596
• @adefazio made their first contribution in #651
• @vz415 made their first contribution in #629
• @mharradon made their first contribution in #696
• @mmhamdy made their first contribution in #682
• @MaanasArora made their first contribution in #688
• @stefanocortinovis made their first contribution in #697
• @Bigpet made their first contribution in #706

Full Changelog: v0.1.7...v0.1.8

Optax 0.1.7

What's Changed

• Remove deprecated field from pyproject.toml, which should hopefully resolve an issue with deploy

Full Changelog: v0.1.6...v0.1.7

Optax 0.1.6

What's Changed

• feat: Add support for is_leaf to tree_map_params, so that tree_map_params works with optax.masked.
• feat: migrate to pyproject.toml for building package by @SauravMaheshkar in #513
• fix(README): small typo in README by @jeertmans in #495
• Add support for PolyLoss (ICLR 2022) by @acforvs in #467

New Contributors

• @jeertmans made their first contribution in #495

Full Changelog: v0.1.5...v0.1.6

Optax 0.1.5

What's Changed

• Fix arXiv link to Optax Optimistic Gradient Descent optimizer by @8bitmp3 in #458
• Fix the Yogi optimizer paper year, change link to NeurIPS site by @8bitmp3 in #461
• Add exponent to cosine decay schedule and warmup + cosine decay by @copybara-service in #476
• Fix typos in docstring by @pomonam in #480
• Fix global_norm() signature by @brentyi in #481
• Fix inject_hyperparams() for python < 3.10. by @copybara-service in #486
• fixed NaN issues in kl_divergence loss function by @LukasMut in #473
• feat(ci/tests): bump setup-python version and enable cache by @SauravMaheshkar in #485
• Better tests for utils by @acforvs in #465
• Run Github CI every day at 03:00. by @copybara-service in #490
• Fix JIT for piecewise_interpolate_schedule, cosine_onecycle_schedule, linear_onecycle_schedule by @brentyi in #504
• Explicitly export "softmax_cross_entropy_with_integer_labels" by @nasyxx in #499
• Add the Lion optimizer, discovered by symbolic program search. by @copybara-service in #500
• Replaces references to jax.numpy.DeviceArray with jax.Array. by @copybara-service in #511
• Update pytypes. by @copybara-service in #514
• Fix pytype failures related to teaching pytype about NumPy scalar types. by @copybara-service in #517
• Release v0.1.5. by @copybara-service in #523

New Contributors

• @pomonam made their first contribution in #480
• @brentyi made their first contribution in #481
• @LukasMut made their first contribution in #473
• @SauravMaheshkar made their first contribution in #485
• @acforvs made their first contribution in #465

Full Changelog: v0.1.4...v0.1.5