Optax Files

What's Changed

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

Full Changelog: https://github.com/google-deepmind/optax/compare/v0.2.1...v0.2.2