Bump the jax group across 1 directory with 5 updates

[dependabot]

Updates the requirements on jax, diffrax, interpax, jax-finufft and quadax to permit the latest version.
Updates jax 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

Updates diffrax to 0.7.2

Release notes

Sourced from diffrax's releases.

Diffrax v0.7.2

• Updated our implicit solvers to work with Optimistix 0.1.0.
• ValueErrors raised from a mistake in the vector field will no longer be masked under a 'terms are not compatible with solver' error.

Full Changelog: patrick-kidger/diffrax@v0.7.1...v0.7.2

Commits

• b3a1885 0.7.2 version bump
• c34f659 Incompatible-with-solver errors no longer masking other underlying ValueErrors
• 33d8c97 update event comments
• f561120 Fix for Optimistix 0.1.0
• 2aa6345 Update optimistix dependency version
• b3b3c8c doc fix
• 6efbea5 0.7.1 version bump
• 4075c88 fix release failing in CI
• dfd555b remove spurious type in union
• 96a905e Added Advanced SDE example and a table of SRKs
• Additional commits viewable in compare view

Updates interpax to 0.3.14

Release notes

Sourced from interpax's releases.

v0.3.14

What's Changed

• Fix reverse mode AD for monotonic splines of constant function
• Bump max jax version to 0.10.0

Full Changelog: f0uriest/interpax@v0.3.13...v0.3.14

Changelog

Sourced from interpax's changelog.

Changelog

• Adds a number of classes that replicate most of the functionality of the corresponding classes from scipy.interpolate :
  • scipy.interpolate.PPoly -> interpax.PPoly
  • scipy.interpolate.Akima1DInterpolator -> interpax.Akima1DInterpolator
  • scipy.interpolate.CubicHermiteSpline -> interpax.CubicHermiteSpline
  • scipy.interpolate.CubicSpline -> interpax.CubicSpline
  • scipy.interpolate.PchipInterpolator -> interpax.PchipInterpolator
• Method "akima" now available for Interpolator.{1D, 2D, 3D} and corresponding functions.
• Method "monotonic" now works in 2D and 3D, where it will preserve monotonicity with respect to each coordinate individually.

v0.2.4

• Fixes for scalar valued query points
• Fixes for interpolating vector valued functions

Full Changelog: f0uriest/interpax@v0.2.3...v0.2.4

v0.2.3

• Add type annotations

Full Changelog: f0uriest/interpax@v0.2.2...v0.2.3

v0.2.2

• Add approx_df to public API

Full Changelog: f0uriest/interpax@v0.2.1...v0.2.2

v0.2.1

• More efficient nearest neighbor search
• Correct slopes for linear interpolation in 2d, 3d
• Fix for cubic2 splines in 2d and 3d Forward and reverse mode AD now fully working and tested

Full Changelog: f0uriest/interpax@v0.2.0...v0.2.1

v0.2.0

... (truncated)

Commits

• 6595ea1 Bump max jax version to 0.10.0 (#175)
• ba31093 Update requirements.txt
• b610c82 build(deps-dev): update pytest-cov requirement from <7.1,>=2.6 to >=2.6,<7.2 ...
• c75d5cf Merge branch 'main' into dependabot/pip/pytest-cov-gte-2.6-and-lt-7.2
• 2a4fcd5 Fix reverse mode AD for monotonic splines of constant function (#171)
• 61b7db8 build(deps-dev): update pytest-cov requirement
• 4b636cc Fix reverse mode AD for monotonic splines of constant function
• e8da61a Update package version constraints in requirements-dev.txt (#170)
• 6a768b1 Update package version constraints in requirements-dev.txt
• 4bdb4a8 Update version constraints for jaxtyping and numpy (#169)
• Additional commits viewable in compare view

Updates jax-finufft to 1.3.1

Release notes

Sourced from jax-finufft's releases.

v1.3.1

This is a patch release to restore wheel compatibility with jax < 0.9.1. The build-time JAX version is now pinned to ensure use of XLA FFI API v0.1. Wheels and conda-forge packages are also now being built for Linux ARM64.

Thanks to @​dfm, @​YigitElma, and @​chrisfinlay for their help this cycle!

The finufft dependency is at 2.5.0.

What's Changed

• build: pin jax==0.8.0 build dependency by @​lgarrison in flatironinstitute/jax-finufft#213
• wheels: add linux-arm64 by @​lgarrison in flatironinstitute/jax-finufft#215

Full Changelog: flatironinstitute/jax-finufft@v1.3.0...v1.3.1

Commits

• 12564d9 wheels: add linux-arm64 (#215)
• 97c5b7c build: pin jax==0.8.0 build dependency (#213)
• 01c802e update finufft to v2.5.0, add CUDA 13 support (#211)
• 7d8266d ops: use p2tz in jax >= 0.9.1
• 86b3655 [pre-commit.ci] pre-commit autoupdate
• 23862a6 build(deps): bump actions/download-artifact from 7 to 8
• 262e323 build(deps): bump actions/upload-artifact from 6 to 7
• aa5df48 [pre-commit.ci] pre-commit autoupdate
• 9d8a870 build(deps): bump pypa/cibuildwheel from 3.3.0 to 3.3.1
• 56f35bb update jvp function for broadcasting support of point grads; add test for bat...
• Additional commits viewable in compare view

Updates quadax to 0.2.13

Release notes

Sourced from quadax's releases.

v0.2.13

What's Changed

• Bump max jax version to 0.10.0

Full Changelog: f0uriest/quadax@v0.2.12...v0.2.13

Commits

• 0da4e28 Bump max jax version to 0.10.0 (#128)
• ea234f3 Ensure arrays are arrays
• 34fe2e7 Update requirements.txt
• 3b223c3 Bump codecov/codecov-action from 5 to 6 (#126)
• f4f0e35 Bump codecov/codecov-action from 5 to 6
• 3d5b332 Update pytest-cov requirement from <7.1,>=2.6 to >=2.6,<7.2 (#125)
• cdbb511 Update pytest-cov requirement from <7.1,>=2.6 to >=2.6,<7.2
• f87334a Update package version constraints in requirements-dev.txt (#124)
• 76cfca4 Update package version constraints in requirements-dev.txt
• 09cc0e7 Update package version constraints in requirements-dev.txt (#118)
• Additional commits viewable in compare view

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    5.54 %    |     3.904e+03      |     4.120e+03      |    216.22    |       26.82        |       23.10        |
  test_proximal_jac_w7x_with_eq_update   |   -0.53 %    |     6.521e+03      |     6.486e+03      |    -34.86    |       117.34       |       119.03       |
  test_proximal_freeb_jac                |    0.76 %    |     1.324e+04      |     1.334e+04      |    100.95    |       63.64        |       61.34        |
  test_proximal_freeb_jac_blocked        |    1.46 %    |     7.579e+03      |     7.689e+03      |    110.32    |       55.62        |       53.90        |
  test_proximal_freeb_jac_batched        |   -0.18 %    |     7.632e+03      |     7.618e+03      |    -13.58    |       55.21        |       54.20        |
  test_proximal_jac_ripple               |    0.28 %    |     3.481e+03      |     3.490e+03      |     9.70     |       41.18        |       40.38        |
  test_proximal_jac_ripple_bounce1d      |    3.30 %    |     3.707e+03      |     3.829e+03      |    122.25    |       51.02        |       50.02        |
  test_eq_solve                          |    2.57 %    |     2.047e+03      |     2.100e+03      |    52.59     |       64.79        |       64.39        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[YigitElma]

I don't understand why Dependabot tries to bump the minimum version. I will manually fix it if the issue is not resolved soon.