Optax 0.1.2

What's Changed

• Improve the documentation and discoverability of set_to_zero. by @copybara-service in #299
• Add pages to documentation for contributors. by @copybara-service in #301
• Replace _build_simple_adam with _build_stateful_sgd. by @n2cholas in #302
• Make masked preserve param structure by @n2cholas in #300
• Allow to set custom schedule for the second momentum in Adafactor. One useful use case is setting the upper bound for the second momentum, as otherwise it converges to 1.0 and effectively freezes the second momentum. It was used in https://arxiv.org/abs/2106.04560. by @copybara-service in #303
• exporting CTC loss for public use by @copybara-service in #305
• Specify that mask must be static in masked wrapper docstring. by @copybara-service in #306
• Clarify the docstring of cosine_decay_schedule. by @copybara-service in #310
• Typo in LARS docstring. by @Olamon in #311
• Fix adam's mu_dtype casting. by @copybara-service in #313
• Fix docs rendering of loss functions. by @grahamgower in #318
• Clarifies optax.adamw(weight_decay) parameter. by @copybara-service in #322
• Enhance ctc_loss to return forward probs. by @yotarok in #321
• Replace jax.tree_multimap, which is now deprecated, with jax.tree_map. by @copybara-service in #330
• Rename argument names for CTC loss functions. by @yotarok in #331

New Contributors

• @Olamon made their first contribution in #311
• @grahamgower made their first contribution in #318
• @yotarok made their first contribution in #321

Full Changelog: v0.1.1...v0.1.2