Bump jax from 0.9.0.1 to 0.10.0

[dependabot]

Bumps jax from 0.9.0.1 to 0.10.0.

Release notes

Sourced from jax's releases.

JAX v0.10.0

• New features:

  • Added ResizeMethod.CUBIC_PYTORCH to jax.image.resize to match PyTorch's bicubic resize (#15768).
  • We now support differentiation of jax.lax.linalg.qr for wide matrices and when full_matrices is True.
  • LAPACK operations are now parallelized along the batch dimension on CPU.
  • Added perturb_singular argument to jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near-zero pivots in the LU decomposition. This is useful for solving numerically singular systems when computing eigenvectors by inverse iteration.
  • jax.scipy.linalg.eigh_tridiagonal now supports computing eigenvectors on CPU and GPU.
  • Added the jax.numpy.ndarray.byteswap method.

• Breaking changes:

  • PartitionSpec objects no longer report themselves to be equal to tuples. Convert tuples to PartitionSpec objects before testing equality.
  • The .vma property has been removed from jax.core.ShapedArray. Use .manual_axis_type.varying instead.
  • JAX CPU devices now report their names as cpu:0, cpu:1, etc. instead of TFRT_CPU_0, TFRT_CPU_1.
  • The config state jax_pmap_shmap_merge has been removed. jax.pmap will now always use the new implementation that wraps jax.jit(jax.shard_map). Please see https://docs.jax.dev/en/latest/migrate_pmap.html for more information.
  • jax.device_put_sharded and jax.device_put_replicated have been removed from the public API and now raise an AttributeError when accessed. Please see https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements for drop-in replacements.
  • The C++ pmap infrastructure has been removed. The following public APIs are no longer available:
    • jax.sharding.PmapSharding
    • From jaxlib.xla_extension: PmapFunction, pmap, NoSharding, Chunked, Unstacked, ShardedAxis, Replicated, ShardingSpec.
    • From jax.interpreters.pxla: MapTracer, PmapExecutable, parallel_callable, shard_args, xla_pmap_p, Chunked, NoSharding, Replicated, ShardedAxis, ShardingSpec, Unstacked, spec_to_indices.
  • The deprecated keyword arguments a, a_min, and a_max to jax.numpy.clip have been removed.
  • Functions jax.numpy.hstack, jax.numpy.vstack, jax.numpy.dstack, jax.numpy.column_stack, jax.numpy.atleast_1d, jax.numpy.atleast_2d, and jax.numpy.atleast_3d no longer accept non-ArrayLike inputs. Doing so previously issued a DeprecationWarning.
  • jax.scipy.stats.rankdata now returns floating point values in all cases, following a similar change in the SciPy 1.18 release.

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.10.0 (April 16, 2026)

• New features:

  • Added ResizeMethod.CUBIC_PYTORCH to {func}jax.image.resize to match PyTorch's bicubic resize ({jax-issue}[#15768](https://github.com/jax-ml/jax/issues/15768)).
  • We now support differentiation of {func}jax.lax.linalg.qr for wide matrices and when full_matrices is True.
  • LAPACK operations are now parallelized along the batch dimension on CPU.
  • Added perturb_singular argument to {func}jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near-zero pivots in the LU decomposition. This is useful for solving numerically singular systems when computing eigenvectors by inverse iteration.
  • {func}jax.scipy.linalg.eigh_tridiagonal now supports computing eigenvectors on CPU and GPU.
  • Added the {meth}jax.numpy.ndarray.byteswap method.

• Breaking changes:

  • PartitionSpec objects no longer report themselves to be equal to tuples. Convert tuples to PartitionSpec objects before testing equality.
  • The .vma property has been removed from jax.core.ShapedArray. Use .manual_axis_type.varying instead.
  • JAX CPU devices now report their names as cpu:0, cpu:1, etc. instead of TFRT_CPU_0, TFRT_CPU_1.
  • The config state jax_pmap_shmap_merge has been removed. jax.pmap will now always use the new implementation that wraps jax.jit(jax.shard_map). Please see https://docs.jax.dev/en/latest/migrate_pmap.html for more information.
  • jax.device_put_sharded and jax.device_put_replicated have been removed from the public API and now raise an AttributeError when accessed. Please see https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements for drop-in replacements.
  • The C++ pmap infrastructure has been removed. The following public APIs are no longer available:
    • jax.sharding.PmapSharding
    • From jaxlib.xla_extension: PmapFunction, pmap, NoSharding, Chunked, Unstacked, ShardedAxis, Replicated, ShardingSpec.
    • From jax.interpreters.pxla: MapTracer, PmapExecutable, parallel_callable, shard_args, xla_pmap_p, Chunked, NoSharding, Replicated, ShardedAxis, ShardingSpec, Unstacked, spec_to_indices.
  • The deprecated keyword arguments a, a_min, and a_max to jax.numpy.clip have been removed.
  • Functions jax.numpy.hstack, jax.numpy.vstack, jax.numpy.dstack, jax.numpy.column_stack, jax.numpy.atleast_1d, jax.numpy.atleast_2d, and jax.numpy.atleast_3d no longer accept non-ArrayLike inputs. Doing so previously issued a DeprecationWarning.
  • {func}jax.scipy.stats.rankdata now returns floating point values in

... (truncated)

Commits

• a33ed61 Add a libtpu version guard in test_dce_sink_prevents_xla_dce
• 79e8ef9 Remove unused EXTRA_INDEX environment variable
• 83e0ac5 Allow shell globbing for libtpu wheel path in pytest_tpu.yml
• 08a228a Prepare for JAX release 0.10.0
• 2362cc7 Fix a bug in WeakValueInterner.
• c12e417 Addressed an old pytype TODO
• 1d0523d Skip a test in OSS that relies on stdout
• 01c0f46 [Pallas] Make VMEM explicit on pallas_call TPU kernel inputs during tracing.
• ad793d6 Use @​immutable for ShapeDtypeStruct and object.setattr in the constructor.
• 40e7c70 Disable the eigh tests under asan.
• 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)