google-deepmind/optax v0.2.5

Optax 0.2.5

Repository: google-deepmind/optax

Tag: v0.2.5

Published: 2025-06-10T16:59:58Z

Prerelease: no

Release notes:

What's Changed

• Extend example gallery entry for linear assignment problem with optimal transport example. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1143
• In flatname space import from the subpackage not from _src. by @copybara-service in https://github.com/google-deepmind/optax/pull/1147
• Fix docs for optax.partition. by @copybara-service in https://github.com/google-deepmind/optax/pull/1150
• Fix complex support for L-BFGS by @gautierronan in https://github.com/google-deepmind/optax/pull/1142
• Add sophia-h optimizer by @evanatyourservice in https://github.com/google-deepmind/optax/pull/979
• Exposing named_chain in docs by @rdyro in https://github.com/google-deepmind/optax/pull/1153
• Split github workflows for lower latency, add ruff by @rdyro in https://github.com/google-deepmind/optax/pull/1156
• Remove longtime deprecated functions. by @copybara-service in https://github.com/google-deepmind/optax/pull/1149
• Fix typo in docstring by @KhawajaAbaid in https://github.com/google-deepmind/optax/pull/1157
• Changes to test.sh and pyproject.toml. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1110
• Minor grammar fix. by @copybara-service in https://github.com/google-deepmind/optax/pull/1158
• Add Muon Optimizer to contrib by @leloykun in https://github.com/google-deepmind/optax/pull/1126
• Fix axis should be int or None in softmax_cross_entropy_with_integer_labels by @rdyro in https://github.com/google-deepmind/optax/pull/1164
• [Triplet Margin Loss] Issue 1118 by @cvnad1 in https://github.com/google-deepmind/optax/pull/1120
• Fix Muon implementation by @leloykun in https://github.com/google-deepmind/optax/pull/1170
• removing scalar_type_of - random.key breaks it by @rdyro in https://github.com/google-deepmind/optax/pull/1169
• Updating optax to use jax.random.key instead of PRNGKey by @copybara-service in https://github.com/google-deepmind/optax/pull/1172
• CI: test against nightly JAX by @jakevdp in https://github.com/google-deepmind/optax/pull/1173
• optax doctests fix by @copybara-service in https://github.com/google-deepmind/optax/pull/1176
• minor change so backtracking init and update return the same state types by @bafflingbits in https://github.com/google-deepmind/optax/pull/1175
• Patching transforms:trace issue due to JAX tree map deprecation by @mustass in https://github.com/google-deepmind/optax/pull/1181
• Adds AdeMAMix Optimizer to contrib by @mathDR in https://github.com/google-deepmind/optax/pull/1104
• Fixes to docs, prelude to docs build CI by @rdyro in https://github.com/google-deepmind/optax/pull/1184
• Add improved version of Hungarian algorithm. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1140
• Add 'plus' option to polyak_sgd. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1180
• Add optax.tree_utils.tree_batch_shape. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1161
• Move configuration for Pylint from .pylintrc to pyproject.toml. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1160
• Fix Typo. by @copybara-service in https://github.com/google-deepmind/optax/pull/1191
• Fix optax CI by relaxing float32 rtol comparison to 1e-6 by @copybara-service in https://github.com/google-deepmind/optax/pull/1194
• Fixing doctests and removing etils.lazy_import for compatibility with Python 3.9 by @copybara-service in https://github.com/google-deepmind/optax/pull/1195
• Add non-negative least squares (NNLS) solver. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1155
• fix: Unexpected tracer in muon optimizer when sharding. by @hlzl in https://github.com/google-deepmind/optax/pull/1193
• Stop gradient in linesearch identity scaling by @younik in https://github.com/google-deepmind/optax/pull/1190
• fix: Correct string formatting for edge labels in linear assignment e… by @k22036 in https://github.com/google-deepmind/optax/pull/1211
• Fixes #1210 by @copybara-service in https://github.com/google-deepmind/optax/pull/1217
• Deprecate multi_transform in favor of partition by @copybara-service in https://github.com/google-deepmind/optax/pull/1216
• Remove unused dependency by @NeilGirdhar in https://github.com/google-deepmind/optax/pull/1219
• Improved docstring for optax.centralize with explanation and example by @shreyans413 in https://github.com/google-deepmind/optax/pull/1220
• Fix doctest error by @copybara-service in https://github.com/google-deepmind/optax/pull/1221
• Ensure optimizers accept extra args by @copybara-service in https://github.com/google-deepmind/optax/pull/1212
• Update _schedule.py by @dpoteryayev in https://github.com/google-deepmind/optax/pull/1222
• allow triggering CI tests manually by @rdyro in https://github.com/google-deepmind/optax/pull/1226
• fixed PartitionState import issue while building html docs by @julurisaichandu in https://github.com/google-deepmind/optax/pull/1232
• Add tree_utils functions to docs that are missing. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1233
• Improve docstring for optax.power_iteration. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1215
• added mathematical description for the ada delta optimizer. by @julurisaichandu in https://github.com/google-deepmind/optax/pull/1235
• Allow learning_rate of None, addressing https://github.com/google-deepmind/optax/issues/1242 by @copybara-service in https://github.com/google-deepmind/optax/pull/1244
• [#1196] renaming tree_scalar_mul to tree_scale and tree_add_scalar_mul to tree_add_scale by @aryanmahawar205 in https://github.com/google-deepmind/optax/pull/1246
• fixing optax CI by @copybara-service in https://github.com/google-deepmind/optax/pull/1251
• Adding pre-commit by @copybara-service in https://github.com/google-deepmind/optax/pull/1250
• Adding pre-commit, uv in test.sh - better caching by @rdyro in https://github.com/google-deepmind/optax/pull/1247
• Fix scaling factor in muon by @stephen-huan in https://github.com/google-deepmind/optax/pull/1249
• Update references to JAX's GitHub repo by @copybara-service in https://github.com/google-deepmind/optax/pull/1252
• Add defaults for min_value and max_value to tree_utils.tree_clip. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1234
• Add how to read lr by…