JAX 0.12.0 Adds an IFRT Proxy Backend and Custom VJP Rematerialization, and Requires Haswell-Class x86-64 CPUs and NumPy 2.2

Overview

JAX, the numerical computing and machine-learning library maintained in the jax-ml GitHub organization, published version 0.12.0. The GitHub release page shows the tag was created on 2026-10-19, and the project’s changelog dates the release October 19, 2026. The release adds new numerical primitives, a way to customize rematerialization for custom gradients, and a backend that connects to a remote IFRT Proxy server. It also removes one API and raises hardware and NumPy requirements.

What Is New

According to the JAX changelog, the new features include:

• jax.scipy.special.erfcinv, the inverse complementary error function.
• jax.lax.polynomial for polynomial evaluation. The changelog says the polynomial primitive is faster and uses less memory when computing gradients.
• jax.custom_vjp.defremat, which customizes how a jax.custom_vjp function is rematerialized when it is differentiated under jax.remat. The changelog says it requires the jax_custom_vjp3 and jax_remat3 implementations. A remat option was also added to jax.custom_gradient, and jax.custom_vjp.defremat_with_logs lets such rules log values from the backward pass.
• A strategy argument on arr.at[...].get() that selects whether the operation lowers to gather, dynamic_slice, or slice.
• A generic_batching option on jax.custom_vjp.defvjp, for batching under jax.vmap when a primal output is shared across the batch.
• Support for 'mxfp4' in jax.nn.get_scaled_dot_general_config, described in the changelog as E2M1 elements with one E8M0 scale per block of 32. According to the changelog, the quantization runs natively on AMD gfx950 GPUs and through a dequantize-and-dot fallback elsewhere.
• 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 a specified mesh axis.
• Generalized Hermitian eigenproblems in jax.scipy.linalg.eigh, jax.profiler.ProfileSession and jax.profiler.export_profile_data for profiling, and the experimental jax.experimental.thunky module for authoring, inspecting, splicing and executing XLA:GPU ThunkExecutor programs from Python.

Proxy backend

The changelog says JAX now registers a proxy backend that connects to an IFRT Proxy server, and names Pathways on Cloud as an example of such a server. To use it, set JAX_PLATFORMS=proxy and set JAX_BACKEND_TARGET to the server address, for example grpc://localhost:29000, per the same changelog.

Breaking Changes and Deprecations

The changelog lists the following breaking changes:

• jax.custom_remat was removed. The replacement is jax.custom_vjp.defremat: the old call custom_remat(f, f_fwd, f_rem, f_bwd) becomes f = jax.custom_vjp(f) followed by f.defremat(f_fwd, f_rem, f_bwd).
• On x86-64, JAX and jaxlib now require an Intel Haswell or newer CPU, which the changelog describes as roughly 2013 onwards and requiring AVX2 and FMA support. Older CPUs that support only AVX are no longer supported, although building from source for them remains possible.
• NumPy 2.1 support was dropped, so JAX now requires NumPy 2.2 or newer.
• jax.numpy.atleast_1d, atleast_2d and atleast_3d now return tuples of arrays rather than lists, matching NumPy 2.0 and later.
• The shape argument of jax.random.multinomial is now the batch shape, so the result has shape (*shape, p.shape[-1]). Code that passed the full output shape should drop its last entry.
• Under jax.vmap, a jax.custom_vjp function whose primal output does not depend on any batched input could previously receive gradients that counted that output’s cotangent once per example. That case now raises an error, and the changelog points to the new generic_batching option as one way to resolve it.
• jax.numpy.cross no longer accepts 2D vector inputs, completing a deprecation that began in JAX 0.11.0. The copy, order and ndmin arguments to jax.numpy.array are now keyword-only.

The changelog also lists jax.lax.psend and jax.lax.precv as deprecated.

Other Changes

Among the changes listed in the changelog, the “web” debugger now defaults to hostname “localhost”, which the changelog calls a safer default that avoids inadvertently opening a port to the world. JAX now uses Bazel 9.2.0 to build from source. Several numerical-accuracy improvements are listed, including for jax.scipy.special.ndtri, jax.numpy.hypot, jax.numpy.sinc and jax.lax.expm1. According to the changelog, the derivative of jax.numpy.hypot at (0, 0) now returns NaN, where it previously returned ±0.5.

What We Don’t Know

The changelog and release page do not say how many users run JAX on AVX-only CPUs, so the practical reach of the Haswell requirement is unclear. Neither source includes benchmarks for the new polynomial primitive beyond the statement that it is faster and uses less memory when computing gradients.