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

Optax 0.1.4.

What's Changed

New features

• Expose adamax and adamaxw by @nasyxx in #433
• Support feeding in extra dictionary data in optax/experimental in #373
• Add NovoGrad optimizer by @DT6A in #385
• added optimistic gradient descent in #387
• Add types to utils.py by @atgctg in #367
• Add hyperparam_dtype override to hyperparam injection in #414
• added typing to linear_algebra.py by @shivance in #413
• Add amsgrad optimizer by @merajhashemi in #382
• Add Hinge Loss by @heytanay in #409

Bug fixes

• [optax] Increase chex version requirement to fix issue #456. in #457
• Start inject_hyperparams step count at 0. in #416
• Add noise before multiplying by the learning rate (noisy_sgd ) by @atgctg in #369

Miscellaneous

• Remove flags from lookahead_mnist example. in #395
• Bring alias docstrings more in line with style guide. in #406
• Bugfix in alias_test: Adamax named but not tested. in #419
• Test that all optimizers can be wrapped in inject_hyperparams. in #420
• Add an example for gradient accumulation. in #425
• [docs] Start adding numerical definitions to key parts of optax. in #430
• [optax] Add basic mnist example based on the lookahead_mnist example. in #436
• Install dependencies for dp-accounting manually. in #375
• Update documentation for AdamW in #376
• [Docs] softmax_cross_entropy_with_integer_label by @lkhphuc in #351
• Update Returns section in gradient transformations' docstrings. in #388
• Update logo and theme for the documentation in #391
• Set the test version of flax used in the equivalence test. in #398
• Add holounic to contributors list. by in #400
• Bring transform docstrings more in line with style guide. by in #405
• Update citation. in #407
• Refine the doc of sigmoid_binary_cross_entropy to not assume the meaning of last dimension. in #418
• fix integer_pow recompilation in _bias_correction by @epignatelli in #329
• Use auto instead of /proc/cpuinfo. in #454

New Contributors

• @atgctg made their first contribution in #367
• @nasyxx made their first contribution in #433
• @shivance made their first contribution in #413
• @epignatelli made their first contribution in #329
• @merajhashemi made their first contribution in #382
• @heytanay made their first contribution in #409

Full Changelog: v0.1.3...v0.1.4