[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.12.0

• New features
  • Added jax.scipy.special.erfcinv for the inverse complementary error function.
  • Added jax.lax.polynomial for polynomial evaluation. The polynomial primitive is faster and uses less memory when computing gradients.
  • Added jax.custom_vjp.defremat, which customizes how a jax.custom_vjp function is rematerialized when it's differentiated under jax.remat. It requires the jax_custom_vjp3 and jax_remat3 implementations.
  • Added a remat option to jax.custom_gradient, for writing jax.custom_vjp.defremat rules with closures.
  • arr.at[...].get() now accepts a strategy argument that allows choosing whether the operation lowers to gather, dynamic_slice, or slice.
  • Added a generic_batching option to jax.custom_vjp.defvjp, for batching under jax.vmap when a primal output is shared across the batch. See jax-301-vmap-shared-outputs.
  • jax.nn.get_scaled_dot_general_config accepts 'mxfp4': E2M1 elements with one E8M0 scale per block of 32, for jax.nn.scaled_dot_general. The quantization is executed natively on AMD gfx950 GPUs and by a dequantize-and-dot fallback elsewhere.
  • JAX now registers a proxy backend, which connects to an IFRT Proxy server such as Pathways on Cloud. To use it, set JAX_PLATFORMS=proxy and set JAX_BACKEND_TARGET to the server address, e.g. grpc://localhost:29000.
  • Added jax.sharding.no_collectives, a context manager that raises an error at trace time in explicit sharding mode if an operation or jax.sharding.reshard introduces a collective across any specified mesh axis.
  • Added jax.custom_vjp.defremat_with_logs, allowing jax.custom_vjp.defremat rules (and jax.custom_gradient with remat=True and with_logs=True) to log values from the backward pass (#40956).
  • jax.scipy.linalg.eigh now supports generalized Hermitian eigenproblems (b != None) for problem types type=1, 2, and 3 (#41162, #41197).
  • Added jax.profiler.ProfileSession and jax.profiler.export_profile_data for session- and context-manager-based profiling.
  • Added jax.experimental.layout.Layout.for_array to query the default device layout for an array.
  • Added a strict_in_out argument to jax.experimental.program_order to place optimization barriers on the first input and/or last output of the scope.
  • Added jax.experimental.thunky for authoring, inspecting, splicing, and executing XLA:GPU ThunkExecutor programs directly from Python.
  • Added JVP rule for jax.lax.nextafter, enabling forward- and reverse-mode autodiff through nextafter (#38679).
  • Widened jax.random.ball's p parameter type from float to RealArray, allowing array-valued norm orders (#41546).

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.12.0 (October 19, 2026)

• New features
  • Added {func}jax.scipy.special.erfcinv for the inverse complementary error function.
  • Added {func}jax.lax.polynomial for polynomial evaluation. The polynomial primitive is faster and uses less memory when computing gradients.
  • Added {func}jax.custom_vjp.defremat, which customizes how a {func}jax.custom_vjp function is rematerialized when it's differentiated under {func}jax.remat. It requires the jax_custom_vjp3 and jax_remat3 implementations.
  • Added a remat option to {func}jax.custom_gradient, for writing {func}jax.custom_vjp.defremat rules with closures.
  • arr.at[...].get() now accepts a strategy argument that allows choosing whether the operation lowers to gather, dynamic_slice, or slice.
  • Added a generic_batching option to {func}jax.custom_vjp.defvjp, for batching under {func}jax.vmap when a primal output is shared across the batch. See {ref}jax-301-vmap-shared-outputs.
  • {func}jax.nn.get_scaled_dot_general_config accepts 'mxfp4': E2M1 elements with one E8M0 scale per block of 32, for {func}jax.nn.scaled_dot_general. The quantization is executed natively on AMD gfx950 GPUs and by a dequantize-and-dot fallback elsewhere.
  • JAX now registers a proxy backend, which connects to an IFRT Proxy server such as Pathways on Cloud. To use it, set JAX_PLATFORMS=proxy and set JAX_BACKEND_TARGET to the server address, e.g. grpc://localhost:29000.
  • Added {func}jax.sharding.no_collectives, a context manager that raises an error at trace time in explicit sharding mode if an operation or {func}jax.sharding.reshard introduces a collective across any specified mesh axis.
  • Added {func}jax.custom_vjp.defremat_with_logs, allowing {func}jax.custom_vjp.defremat rules (and {func}jax.custom_gradient with remat=True and with_logs=True) to log values from the backward pass ({jax-issue}[#40956](https://github.com/jax-ml/jax/issues/40956)).
  • {func}jax.scipy.linalg.eigh now supports generalized Hermitian eigenproblems (b != None) for problem types type=1, 2, and 3 ({jax-issue}[#41162](https://github.com/jax-ml/jax/issues/41162), {jax-issue}[#41197](https://github.com/jax-ml/jax/issues/41197)).
  • Added {class}jax.profiler.ProfileSession and {func}jax.profiler.export_profile_data for session- and context-manager-based profiling.
  • Added {meth}jax.experimental.layout.Layout.for_array to query the default device layout for an array.
  • Added a strict_in_out argument to {func}jax.experimental.program_order to place optimization barriers on the first input and/or last output of the scope.
  • Added {mod}jax.experimental.thunky for authoring, inspecting, splicing, and executing XLA:GPU ThunkExecutor programs directly from Python.
  • Added JVP rule for {func}jax.lax.nextafter, enabling forward- and reverse-mode autodiff through nextafter ({jax-issue}[#38679](https://github.com/jax-ml/jax/issues/38679)).
  • Widened {func}jax.random.ball's p parameter type from float to

... (truncated)

Commits

• 8dcb7ae Prepare for JAX release 0.12.0
• ca20b9c [ROCm][Triton] Fix race on the HSACO file in GetModuleImage.
• 3c0096f Fix unbound bazel_startup_options array expansion under set -u on macOS B...
• b4bf3dd Update XLA dependency to revision d867969800d4a8fd045505eabc65e1e66f30794f.
• 140e22b Tag gpu_torch_test and mgpu_torch_test with noassert.
• 94b5121 Disable svd_test on CPU under AddressSanitizer (noasan) to avoid timeouts.
• 9734e16 [mgpu] Add pass to lower Mosaic GPU barriers to i64.
• 03c93f0 [jax] Fix reverse-mode AD for csr_matvec and csr_matmat with transpose=True
• 8f51bfd Skip ShardMapTest.test_unreduced_max_min_nested_shard_map_add_manual_axes w...
• cf584e9 [pallas:mgpu] Surfaced im2col mode support in copy_gmem_to_smem
• 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)

[dependabot]

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.