Optax Files

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 @copybara-service in https://github.com/google-deepmind/optax/pull/1253
• simplify centralize by @copybara-service in https://github.com/google-deepmind/optax/pull/1254
• type stability of state in optimizers, tested via jax.lax.cond by @copybara-service in https://github.com/google-deepmind/optax/pull/1218
• Add Simplified AdEMAMix. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1206
• Change type of num_linesearch_steps for backtracking linesearch by @copybara-service in https://github.com/google-deepmind/optax/pull/1258
• Fix nnls to handle batch shapes correctly. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1200
• Edit sigmoid_binary_cross_entropy and sigmoid_focal_loss to work correctly with non-array label arguments. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1261
• Minor improvements to tree_batch_shape. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1260
• Fix docs for optax.contrib.dog. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1262
• Improve docstrings for schedule functions in optax/schedules/_schedule.py by @DakshBegani in https://github.com/google-deepmind/optax/pull/1255
• Fix docs building by @copybara-service in https://github.com/google-deepmind/optax/pull/1263
• Note that default eps values of Adabelief don't match paper. by @copybara-service in https://github.com/google-deepmind/optax/pull/1264
• Seed key v3 by @Tomas542 in https://github.com/google-deepmind/optax/pull/1240
• Abandon python 3.9, test on 3.12 by @copybara-service in https://github.com/google-deepmind/optax/pull/1266
• Clean dependencies by @copybara-service in https://github.com/google-deepmind/optax/pull/1270
• Muon: Add weight decay by @leloykun in https://github.com/google-deepmind/optax/pull/1275
• Fix MuonState docstring by @EIFY in https://github.com/google-deepmind/optax/pull/1276
• Fix duplicate-bases error upon running test.sh by directing pylint to run on optax directory only. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1278
• add freezing parameters examples [Addresses Issue [#296]] by @pranavagrawaI in https://github.com/google-deepmind/optax/pull/1274
• Implementing ADOPT in Optax by @tinker495 in https://github.com/google-deepmind/optax/pull/1257
• Automated Code Change by @copybara-service in https://github.com/google-deepmind/optax/pull/1281
• Update JAX nightly index usage by @copybara-service in https://github.com/google-deepmind/optax/pull/1282
• Configure tests to explicitly use jax_threefry_partitionable=False. by @copybara-service in https://github.com/google-deepmind/optax/pull/1289
• Typos in reduce_on_plateau.ipynb by @ddrous in https://github.com/google-deepmind/optax/pull/1293
• Add license check for all Python files to pre-commit by @copybara-service in https://github.com/google-deepmind/optax/pull/1296
• Remove clutter from MNIST examples by @Balint-H in https://github.com/google-deepmind/optax/pull/1285
• Add Freezing Utilities: freeze, selective_optimizer, and Prefix-Mask Support by @pranavagrawaI in https://github.com/google-deepmind/optax/pull/1284
• Unify tree_(l1|l2|linf)_norm into tree_norm & introduce optax.tree.* aliases. by @copybara-service in https://github.com/google-deepmind/optax/pull/1299
• update optax getting started colab link by @copybara-service in https://github.com/google-deepmind/optax/pull/1301
• update freezing parameter examples with new utilities freeze and selective_transform by @pranavagrawaI in https://github.com/google-deepmind/optax/pull/1300
• Docs: Reorganize warmup schedules in API (fixes [#835]) by @Kenxpx in https://github.com/google-deepmind/optax/pull/1303
• Fix for seed-> key in optax (PR#1240 was not backward compatible) by @copybara-service in https://github.com/google-deepmind/optax/pull/1305
• remove trailing whitespace in pre-commit and disallow CR in line endings by @copybara-service in https://github.com/google-deepmind/optax/pull/1307
• upgrade setup-python version in github actions by @copybara-service in https://github.com/google-deepmind/optax/pull/1311
• Speed up optax github CI by @copybara-service in https://github.com/google-deepmind/optax/pull/1312
• Remove idempotent wording from GradientTransformation docs. by @copybara-service in https://github.com/google-deepmind/optax/pull/1315
• Fix many issues with make_perturbed_fun by @adconner in https://github.com/google-deepmind/optax/pull/1310
• Doctest fix by @copybara-service in https://github.com/google-deepmind/optax/pull/1316
• Minor cleanup for optax.projections. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1277
• Fix typo in docstring for simplified_ademamix. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1322
• Reorder entries of perturbations.rst to put distribution classes together. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1317
• Update version number for release. by @copybara-service in https://github.com/google-deepmind/optax/pull/1326

New Contributors

• @gautierronan made their first contribution in https://github.com/google-deepmind/optax/pull/1142
• @evanatyourservice made their first contribution in https://github.com/google-deepmind/optax/pull/979
• @rdyro made their first contribution in https://github.com/google-deepmind/optax/pull/1153
• @KhawajaAbaid made their first contribution in https://github.com/google-deepmind/optax/pull/1157
• @leloykun made their first contribution in https://github.com/google-deepmind/optax/pull/1126
• @cvnad1 made their first contribution in https://github.com/google-deepmind/optax/pull/1120
• @bafflingbits made their first contribution in https://github.com/google-deepmind/optax/pull/1175
• @mustass made their first contribution in https://github.com/google-deepmind/optax/pull/1181
• @mathDR made their first contribution in https://github.com/google-deepmind/optax/pull/1104
• @younik made their first contribution in https://github.com/google-deepmind/optax/pull/1190
• @k22036 made their first contribution in https://github.com/google-deepmind/optax/pull/1211
• @shreyans413 made their first contribution in https://github.com/google-deepmind/optax/pull/1220
• @dpoteryayev made their first contribution in https://github.com/google-deepmind/optax/pull/1222
• @julurisaichandu made their first contribution in https://github.com/google-deepmind/optax/pull/1232
• @aryanmahawar205 made their first contribution in https://github.com/google-deepmind/optax/pull/1246
• @stephen-huan made their first contribution in https://github.com/google-deepmind/optax/pull/1249
• @DakshBegani made their first contribution in https://github.com/google-deepmind/optax/pull/1255
• @Tomas542 made their first contribution in https://github.com/google-deepmind/optax/pull/1240
• @EIFY made their first contribution in https://github.com/google-deepmind/optax/pull/1276
• @pranavagrawaI made their first contribution in https://github.com/google-deepmind/optax/pull/1274
• @tinker495 made their first contribution in https://github.com/google-deepmind/optax/pull/1257
• @ddrous made their first contribution in https://github.com/google-deepmind/optax/pull/1293
• @Balint-H made their first contribution in https://github.com/google-deepmind/optax/pull/1285
• @Kenxpx made their first contribution in https://github.com/google-deepmind/optax/pull/1303
• @adconner made their first contribution in https://github.com/google-deepmind/optax/pull/1310

Full Changelog: https://github.com/google-deepmind/optax/compare/v0.2.4...v0.2.5