google-deepmind/optax v0.2.7

Optax 0.2.7

Repository: google-deepmind/optax

Tag: v0.2.7

Published: 2026-02-05T20:48:40Z

Prerelease: no

Release notes:

What's Changed

• Update Optax version to 0.2.7.dev. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1420
• Fix doctest by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1424
• Fix piecewise interpolate with 0 first split by @rdyro in https://github.com/google-deepmind/optax/pull/1425
• Expose weight decay as a schedule option in all alias optimizers. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1427
• Skip schedule free tests casting from complex to float. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1428
• Add monitor and measure_with_ema to Optax transformations. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1430
• Add consistent rms scaling for muon update by @shuningjin in https://github.com/google-deepmind/optax/pull/1435
• Add microbatch transformation to optax/experimental. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1434
• Add experimental aggregators in optax. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1436
• Add missing exports to optax/__init__.py. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1433
• Adding microbatch to the docs by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1443
• Add a few sharding-related tests to optax. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1450
• Allow instantiating optimizers before jax has been initialized. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1454
• Adding gradient variance tracking. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1451
• Resolve remaining sharding test failures in optax. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1457
• Add optional "in_axes" and "argnames" kwargs to microbatch, which will allow for natural composition with jax.vmap. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1442
• Honor inject_hyperparams dtypes if passed as jax.Arrays. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1460
• Update reshape_batch_axis to include sharding information when in explicit sharding mode. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1459
• Fix typo in L-BFGS debug information section by @partev in https://github.com/google-deepmind/optax/pull/1465
• Deprecate second order utilities. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1466
• Deprecate optax.global_norm in favor of optax.tree.norm. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1368
• Add signum optimizer by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1463
• Use internal warn_deprecated_function instead of the Chex version. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1469
• Define internal assert_trees_all_{close, equal} functions and use them in tests instead of the Chex versions. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1481
• Replace (non-test) Chex assertions, usually with ValueErrors. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1483
• Replace chex.Numeric with jax.typing.ArrayLike. Replace chex.Scalar with float/int as appropriate. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1479
• Update readme by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1496
• Enforce keyword-only arguments for optional parameters in Optax losses. Disable or fix existing Pytype bugs that surfaced as a result of this change. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1505
• Remove use of cast_tree in favor of optax.tree.cast by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1477
• Minor doc fixes: Include hyperlinks to functions, classes references and correct typos by @rajasekharporeddy in https://github.com/google-deepmind/optax/pull/1511
• [optax] Remove jax_pmap_shmap_merge=False flag in contrib/_complex_valued. #jax-fixit by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1515
• Enforce keyword-only arguments for optional parameters in Optax losses. Disable or fix existing Pytype bugs that surfaced as a result of this change. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1516
• Fix argument order in scale_by_distance_over_gradients by @Aaryan-549 in https://github.com/google-deepmind/optax/pull/1501
• Fix momo crash when loss value is a Python float by @Aaryan-549 in https://github.com/google-deepmind/optax/pull/1502
• Clarify difference between kl_divergence and convex_kl_divergence by @zer-art in https://github.com/google-deepmind/optax/pull/1514
• Fix broken test in numerics_test. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1522
• Add guidelines on contributing AI generated code, the same as JAX's. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1526
• Move _microbatching.py to microbatching.py and update __init__.py so we can directly import microbatching without causing pytype errors. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1525
• Fix up microbatching documentation. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1527
• Fix conjugation in Newton-Schulz iterator and update tests for comple… by @eirikfagerbakke in https://github.com/google-deepmind/optax/pull/1506
• Fix microbatching.Accumulator in docs. by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1529
• Upgrade GitHub Actions to latest versions by @salmanmkc in https://github.com/google-deepmind/optax/pull/1534
• Add axis and where arguments to smooth_labels function. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1492
• Fix _projections_test.py: Remove prints and use optax.tree.allclose. by @carlosgmartin in https://github.com/google-deepmind/optax/pull/1439
• updating github actions versions by @copybara-service[bot] in https://github.com/google-deepmind/optax/pull/1540
• Fix typo in alias.py documentation by @partev…