[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.12.0.1

• Bug fixes
  • Built C++-generated intrinsic LLVM IR on Windows with clang-cl instead of emitting empty bitcode, enabling vectorized CPU intrinsics on Windows (openxla/xla@3af77b5).

Changelog

Sourced from jax's changelog.

JAX 0.12.0.1 (October 23, 2026)

• Bug fixes
  • Built C++-generated intrinsic LLVM IR on Windows with clang-cl instead of emitting empty bitcode, enabling vectorized CPU intrinsics on Windows (openxla/xla@3af77b5).

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

... (truncated)

Commits

• dc4af68 Prepare for JAX release 0.12.0.1
• 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...
• Additional commits viewable in compare view