build(deps): bump optax from 0.2.5 to 0.2.8

[dependabot]

Bumps optax from 0.2.5 to 0.2.8.

Release notes

Sourced from optax's releases.

Optax 0.2.8

What's Changed

• Following the JAX 0.9.2 release, the jax_pmap_shmap_merge config flag was removed so that the jax.pmap implementation is always based on jax.jit and jax.shard_map, and opting into the old jax.pmap behavior is no longer an option. Optax had opted into the old behavior to give users time to migrate, and as of Optax 0.2.8 this is no longer supported. This changed shouldn't impact most users, but if you experience errors or performance regressions as a result of it, you can update your code following JAX's migration guide (or use JAX 0.9.2 or earlier and set jax.config.update("jax_pmap_shmap_merge", False)).
• Explicitly specify the dtype of the gradient accumulator in the MultiStep transform. by @​copybara-service[bot] in google-deepmind/optax#1605
• feat: add preconditioning and coef presets to muon by @​massena-t in google-deepmind/optax#1602
• Backwards compatibility export for the newton schulz iterator by @​copybara-service[bot] in google-deepmind/optax#1608
• Remove TensorFlow dependency in Adversarial training example by @​rajasekharporeddy in google-deepmind/optax#1609
• Improve lookahead docstrings with example and usage notes by @​rdyro in google-deepmind/optax#1619
• Make sure inject_hyperparams uses the dtype inferred from parameters... by @​copybara-service[bot] in google-deepmind/optax#1615
• Memory-optimization for microbatching. by @​copybara-service[bot] in google-deepmind/optax#1623
• Remove TensorFlow dependency and migrate mlp_mnist to Flax NNX by @​selamw1 in google-deepmind/optax#1536
• Let inject use the highest dtype found in the params as the default dtype of params. by @​copybara-service[bot] in google-deepmind/optax#1628
• Support scheduling alpha for AdEMAmix by @​copybara-service[bot] in google-deepmind/optax#1630
• [JAX] Suppress type errors found by pytype after correcting definition of jax.typing.ArrayLike. by @​copybara-service[bot] in google-deepmind/optax#1629
• [JAX] Suppress type errors found by pytype after correcting definition of jax.typing.ArrayLike. by @​copybara-service[bot] in google-deepmind/optax#1633

New Contributors

• @​massena-t made their first contribution in google-deepmind/optax#1602
• @​selamw1 made their first contribution in google-deepmind/optax#1536

Full Changelog: google-deepmind/optax@v0.2.7...v0.2.8

Optax 0.2.7

What's Changed

• Update Optax version to 0.2.7.dev. by @​copybara-service[bot] in google-deepmind/optax#1420
• Fix doctest by @​copybara-service[bot] in google-deepmind/optax#1424
• Fix piecewise interpolate with 0 first split by @​rdyro in google-deepmind/optax#1425
• Expose weight decay as a schedule option in all alias optimizers. by @​copybara-service[bot] in google-deepmind/optax#1427
• Skip schedule free tests casting from complex to float. by @​copybara-service[bot] in google-deepmind/optax#1428
• Add monitor and measure_with_ema to Optax transformations. by @​copybara-service[bot] in google-deepmind/optax#1430
• Add consistent rms scaling for muon update by @​shuningjin in google-deepmind/optax#1435
• Add microbatch transformation to optax/experimental. by @​copybara-service[bot] in google-deepmind/optax#1434
• Add experimental aggregators in optax. by @​copybara-service[bot] in google-deepmind/optax#1436
• Add missing exports to optax/init.py. by @​carlosgmartin in google-deepmind/optax#1433
• Adding microbatch to the docs by @​copybara-service[bot] in google-deepmind/optax#1443
• Add a few sharding-related tests to optax. by @​copybara-service[bot] in google-deepmind/optax#1450
• Allow instantiating optimizers before jax has been initialized. by @​copybara-service[bot] in google-deepmind/optax#1454
• Adding gradient variance tracking. by @​copybara-service[bot] in google-deepmind/optax#1451
• Resolve remaining sharding test failures in optax. by @​copybara-service[bot] in google-deepmind/optax#1457
• Add optional "in_axes" and "argnames" kwargs to microbatch, which will allow for natural composition with jax.vmap. by @​copybara-service[bot] in google-deepmind/optax#1442
• Honor inject_hyperparams dtypes if passed as jax.Arrays. by @​copybara-service[bot] in google-deepmind/optax#1460
• Update reshape_batch_axis to include sharding information when in explicit sharding mode. by @​copybara-service[bot] in google-deepmind/optax#1459
• Fix typo in L-BFGS debug information section by @​partev in google-deepmind/optax#1465
• Deprecate second order utilities. by @​copybara-service[bot] in google-deepmind/optax#1466
• Deprecate optax.global_norm in favor of optax.tree.norm. by @​carlosgmartin in google-deepmind/optax#1368
• Add signum optimizer by @​copybara-service[bot] in google-deepmind/optax#1463
• Use internal warn_deprecated_function instead of the Chex version. by @​copybara-service[bot] in google-deepmind/optax#1469
• Define internal assert_trees_all_{close, equal} functions and use them in tests instead of the Chex versions. by @​copybara-service[bot] in google-deepmind/optax#1481
• Replace (non-test) Chex assertions, usually with ValueErrors. by @​copybara-service[bot] in google-deepmind/optax#1483
• Replace chex.Numeric with jax.typing.ArrayLike. Replace chex.Scalar with float/int as appropriate. by @​copybara-service[bot] in google-deepmind/optax#1479

... (truncated)

Commits

• 107e9ad Release Optax version 0.2.8.
• 70df570 [JAX] Suppress type errors found by pytype after correcting definition of jax...
• d015bc4 [JAX] Suppress type errors found by pytype after correcting definition of jax...
• 82601b4 Support scheduling alpha for AdEMAmix
• a6ba37b Let inject use the highest dtype found in the params as the default dtype of ...
• 4a53ba3 Merge pull request #1536 from selamw1:mnl_nnx
• b24c5fa Merge branch 'google-deepmind:main' into mnl_nnx
• 9ff5961 Memory-optimization for microbatching.
• 1e52e8c Make sure inject_hyperparams uses the dtype inferred from parameters...
• a6c2942 Merge pull request #1619 from rdyro:lookahead-imprv
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)