deps(deps): bump the jax-ecosystem group across 1 directory with 5 updates

[dependabot]

Bumps the jax-ecosystem group with 5 updates in the / directory:

Package	From	To
jax	0.9.0	0.10.2
jaxlib	0.9.0	0.10.2
optax	0.2.6	0.2.8
jaxtyping	0.3.6	0.3.11
equinox	0.13.5	0.13.8

Updates jax from 0.9.0 to 0.10.2

Release notes

Sourced from jax's releases.

JAX v0.10.2

• New features
  • Added jax.scipy.linalg.invhilbert for the closed-form inverse of the Hilbert matrix ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added jax.scipy.linalg.invpascal for the inverse of the Pascal matrix ([#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added jax.scipy.linalg.fiedler_companion for constructing the pentadiagonal Fiedler companion matrix of a polynomial ([#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added jax.ShapeDtypeStruct.like -- a shortcut for constructing a jax.ShapeDtypeStruct from an object with shape and dtype attributes.
• We have made a number of improvements to JAX tracing performance.

JAX v0.10.1

• New features

  • Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098).
  • Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144).
  • Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144).
  • Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144).
  • Added jax.scipy.linalg.leslie for constructing Leslie matrices (#10144).
  • Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients (#10144).
  • Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices (#10144).
  • Added jax.scipy.linalg.helmert for constructing Helmert matrices (#10144).
  • Moved RNG APIs from "implementations" to dtypes (#27854):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.10.2 (June 17, 2026)

• New features
  • Added {func}jax.scipy.linalg.invhilbert for the closed-form inverse of the Hilbert matrix ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.invpascal for the inverse of the Pascal matrix ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler_companion for constructing the pentadiagonal Fiedler companion matrix of a polynomial ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.ShapeDtypeStruct.like -- a shortcut for constructing a {class}jax.ShapeDtypeStruct from an object with shape and dtype attributes.

JAX 0.10.1 (May 20, 2026)

• New features

  • Added ResizeMethod.AREA to {func}jax.image.resize, which matches TensorFlow's AREA resizing ({jax-issue}[#20098](https://github.com/jax-ml/jax/issues/20098)).
  • Added {func}jax.scipy.linalg.hadamard for constructing Hadamard matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.circulant for constructing circulant matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.dft for constructing discrete Fourier transform matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.leslie for constructing Leslie matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.helmert for constructing Helmert matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.special.boxcox and {func}jax.scipy.special.boxcox1p for the Box-Cox power transformation.
  • Moved RNG APIs from "implementations" to dtypes ({jax-issue}[#27854](https://github.com/jax-ml/jax/issues/27854)):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to {func}jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future

... (truncated)

Commits

• 990e6a0 Disable gpu_pallas_distributed_nvshmem_test_gpu which was failing at the
• 4437b23 Synchronize PGLE profiler tests before stopping traces
• 3815e4e Add a workaround for linux aarch64 python 3.14t pytest builds.
• 272e378 Refactor Windows Docker GCE metadata access
• 310278e Prepare for JAX release 0.10.2
• 7db3a70 [Pallas:MGPU] Add support for batch dimensions to async_store_tmem.
• 5597e90 [Mosaic:GPU] Disable set_xla_gpu_enable_allocator_spatial_partitioning for nv...
• 728e4bf Add functionality to blanket disable the custom call safety checks in export.
• 26113fa Merge pull request #38432 from vfdev-5:pallas-emit-pipeline-raises-error-empt...
• c15bacc Pallas emit_pipeline raises ValueError on empty grid
• Additional commits viewable in compare view

Updates jaxlib from 0.9.0 to 0.10.2

Changelog

Sourced from jaxlib's changelog.

JAX 0.10.2 (June 17, 2026)

• New features
  • Added {func}jax.scipy.linalg.invhilbert for the closed-form inverse of the Hilbert matrix ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.invpascal for the inverse of the Pascal matrix ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler_companion for constructing the pentadiagonal Fiedler companion matrix of a polynomial ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.ShapeDtypeStruct.like -- a shortcut for constructing a {class}jax.ShapeDtypeStruct from an object with shape and dtype attributes.

JAX 0.10.1 (May 20, 2026)

• New features

  • Added ResizeMethod.AREA to {func}jax.image.resize, which matches TensorFlow's AREA resizing ({jax-issue}[#20098](https://github.com/jax-ml/jax/issues/20098)).
  • Added {func}jax.scipy.linalg.hadamard for constructing Hadamard matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.circulant for constructing circulant matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.dft for constructing discrete Fourier transform matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.leslie for constructing Leslie matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.helmert for constructing Helmert matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.special.boxcox and {func}jax.scipy.special.boxcox1p for the Box-Cox power transformation.
  • Moved RNG APIs from "implementations" to dtypes ({jax-issue}[#27854](https://github.com/jax-ml/jax/issues/27854)):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to {func}jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future

... (truncated)

Commits

• 990e6a0 Disable gpu_pallas_distributed_nvshmem_test_gpu which was failing at the
• 4437b23 Synchronize PGLE profiler tests before stopping traces
• 3815e4e Add a workaround for linux aarch64 python 3.14t pytest builds.
• 272e378 Refactor Windows Docker GCE metadata access
• 310278e Prepare for JAX release 0.10.2
• 7db3a70 [Pallas:MGPU] Add support for batch dimensions to async_store_tmem.
• 5597e90 [Mosaic:GPU] Disable set_xla_gpu_enable_allocator_spatial_partitioning for nv...
• 728e4bf Add functionality to blanket disable the custom call safety checks in export.
• 26113fa Merge pull request #38432 from vfdev-5:pallas-emit-pipeline-raises-error-empt...
• c15bacc Pallas emit_pipeline raises ValueError on empty grid
• Additional commits viewable in compare view

Updates optax from 0.2.6 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

Updates jaxtyping from 0.3.6 to 0.3.11

Release notes

Sourced from jaxtyping's releases.

jaxtyping v0.3.11

• Bugfix: the import hook now correctly handles the @no_type_check decorator. (Thanks @​jeertmans! patrick-kidger/jaxtyping#393)

Full Changelog: patrick-kidger/jaxtyping@v0.3.10...v0.3.11

jaxtyping v0.3.10

Bugfixes

• add u32[4] as a possible PRNGKeyArray type (Thanks @​jondeaton! patrick-kidger/jaxtyping#388)
• Make variadic objects (e.g. Float[Array, "..."]) survive cloudpickle round-trip. (Thanks @​hmgaudecker! patrick-kidger/jaxtyping#391)

Infrastructure

• Use dependency groups, uv, and local pre-commit hooks (patrick-kidger/jaxtyping#386)
• Add pyrefly as a type checker in pre-commit and dev dependencies (Thanks @​lolpack! patrick-kidger/jaxtyping#387)

New Contributors

• @​patrick-kidger-bot made their first contribution in patrick-kidger/jaxtyping#385
• @​lolpack made their first contribution in patrick-kidger/jaxtyping#387
• @​jondeaton made their first contribution in patrick-kidger/jaxtyping#388
• @​hmgaudecker made their first contribution in patrick-kidger/jaxtyping#391

Full Changelog: patrick-kidger/jaxtyping@v0.3.9...v0.3.10

jaxtyping v0.3.9

Bugfix: typechecking of pytrees of |-style unions, e.g. isinstance(..., PyTree[int | bool]), will no longer silently always pass. patrick-kidger/jaxtyping#381

Full Changelog: patrick-kidger/jaxtyping@v0.3.8...v0.3.9

jaxtyping v0.3.8

Another bugfix release for numpy.typing.ArrayLike on numpy 2.4.0, I think this time only needed Python 3.13+. #380

Full Changelog: patrick-kidger/jaxtyping@v0.3.7...v0.3.8

jaxtyping v0.3.7

Bugfix: SomeDtype[Union[numpy.typing.ArrayLike, ...], ...] will no longer crash in numpy 2.4.0. (#374)

Full Changelog: patrick-kidger/jaxtyping@v0.3.6...v0.3.7

Commits

• f4083b2 v0.3.11
• c80c7ad fix(lib): import_hook now correctly handles @no_type_check decorator
• fca7077 Make _array_types sentinels survive cloudpickle round-trip.
• a3e7bef add u32[4] as possible PRNGKeyArray type (#388)
• 1e79d13 Add pyrefly as a type checker in pre-commit and dev dependencies (#387)
• 1d54108 Use dependency groups, uv, and local pre-commit hooks (#386)
• 0609755 bump pymdown-extensions; pin pygments
• f7e0088 0.3.9
• aae2fa7 Fix typechecking of pytree of new-style union.
• f6477d4 bump release pipeline to 3.12 to avoid jax weirdness on 3.11
• Additional commits viewable in compare view

Updates equinox from 0.13.5 to 0.13.8

Release notes

Sourced from equinox's releases.

Equinox v0.13.8

This is a minor release, largely for internal maintenance purposes.

• Update awesome-list.md by @​LambdaP in patrick-kidger/equinox#1225
• Add return type to BoundMethod by @​nstarman in patrick-kidger/equinox#1222
• Added progress meters to equinox.internal (=not a public API) (Thanks @​ASKabalan! patrick-kidger/equinox#1226)

New Contributors

• @​LambdaP made their first contribution in patrick-kidger/equinox#1225
• @​ASKabalan made their first contribution in patrick-kidger/equinox#1226

Full Changelog: patrick-kidger/equinox@v0.13.7...v0.13.8

Equinox v0.13.7

Compatibility:

• JAX v0.10.0 deprecated jax.core.find_top_trace; now using jax.extend.core.find_top_trace when available. (#1218)
• JAX v0.10.0 removed jax.core.get_aval, jax.core.mapped_aval, jax.core.unmapped_aval; suppressed the resulting type errors. (#1218)
• JAX v0.9.2 now compatible with equinox.internal.finalise_jaxpr. (Thanks @​kasper0406! #1200)

Other:

• Module.__new__ now returns Self rather than Module. (Thanks @​RunnersNum40! #1207)
• Switched to uv and dependency groups for development. (#1218)

New awesome-list entries:

• Add Parax and ParamRF to awesome-list (Thanks @​gvcallen! #1209)
• add peex to awesome-list (Thanks @​sammccallum! #1208)
• Add GPJax to awesome-list (Thanks @​thomaspinder! #1213)
• Add chronos2 implementation to awesome-list (Thanks @​theorashid! #1214)

New Contributors

• @​kasper0406 made their first contribution in patrick-kidger/equinox#1200
• @​patrick-kidger-bot made their first contribution in patrick-kidger/equinox#1205
• @​gvcallen made their first contribution in patrick-kidger/equinox#1209
• @​sammccallum made their first contribution in patrick-kidger/equinox#1208
• @​thomaspinder made their first contribution in patrick-kidger/equinox#1213
• @​theorashid made their first contribution in patrick-kidger/equinox#1214

Full Changelog: patrick-kidger/equinox@v0.13.6...v0.13.7

Equinox v0.13.6

Bug fix: overriding an AbstractClassVar with a concrete instance field will no longer cause a crash when attempting to instantiate the module. (#1198)

Full Changelog: patrick-kidger/equinox@v0.13.5...v0.13.6

Commits

• dfadb04 0.13.8 release
• d6ae901 Add progress meter to equinox
• 87f4072 Add return type to BoundMethod
• 53b3c75 Update awesome-list.md
• 5972a41 Bump version to 0.13.7
• d77d057 Use dependency groups, uv, and local pre-commit hooks (#1218)
• 7f7b46c Update awesome-list.md
• 9e7cc3e Add GPJax to list (#1213)
• e41cdac Change Module.__new__ return self type rather than Module (#1207)
• 940f903 add peex to awesome-list
• Additional commits viewable in compare view

[dependabot]

Labels

The following labels could not be found: automated. Please create it before Dependabot can add it to a pull request.

Please fix the above issues or remove invalid values from dependabot.yml.

[thomaspinder]

Superseded by #734, which groups all Python updates into a single PR (config change in #732).

[dependabot]

This pull request was built based on a group rule. Closing it will not ignore any of these versions in future pull requests.

To ignore these dependencies, configure ignore rules in dependabot.yml