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