google-deepmind/optax v0.1.2 on GitHub

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

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