pypi / jax
jax
What changed in jax 0.11 to 0.12, from its official release notes, upgrade guides and advisories. Sources checked 2026-10-30.
30 breaking2 security4 deprecations30 new APIs and options6 performance4 behavior and defaults3 tooling56 fixes and other
See which of these apply to your code, ranked, in your agent:
npx @unstale/cli initjax 0.1269 changes
Breaking changes 17
Removed; use jax.custom_vjp.defremat instead.
Under jax.vmap, shared primal outputs that do not depend on batched inputs now raise an error instead of potentially counting their cotangent once per example; generic_batching can resolve the case.
Now clears the caller's concrete and abstract mesh when invoking host callbacks, matching jax.debug.callback.
Now returns tuples of arrays rather than lists, matching NumPy 2.0+.
Now returns tuples of arrays rather than lists, matching NumPy 2.0+.
Now returns tuples of arrays rather than lists, matching NumPy 2.0+.
Scalar constructors now have proper type annotations instead of Any, potentially producing typecheck errors where they were misused.
Now bit-exact for normal floating-point numbers using bitwise operations; subnormal inputs and underflowing outputs are flushed to signed zero.
Improved numerical accuracy and derivatives using power-of-two scaling and a custom JVP rule; the derivative at (0, 0) now returns NaN instead of ±0.5.
Scalar constructors now have proper type annotations instead of Any, potentially producing typecheck errors where they were misused.
Now bit-exact for normal floating-point numbers using bitwise operations; subnormal inputs and underflowing outputs are flushed to signed zero.
Now raises AttributeError when unknown attributes are assigned.
The shape argument now means batch shape, and results have shape (*shape, p.shape[-1]); callers passing the full output shape should drop its last entry.
Updated the minimum pinned libtpu version to 0.0.50.
Dropped support for NumPy 2.1; JAX now requires NumPy 2.2 or newer.
jax.checkpoint, jax.custom_vjp, and jax.shard_map now error if the same Ref is passed more than once or both closed over and passed as an argument.
JAX and jaxlib now require an Intel Haswell or newer CPU with AVX2 and FMA; older AVX-only CPUs are unsupported, though source builds remain possible.
Security 1
Now defaults to hostname localhost, avoiding inadvertently opening a port to the world.
Deprecations 2
Deprecated.
Deprecated.
New APIs and options 20
The new strategy argument chooses whether the operation lowers to gather, dynamic_slice, or slice.
Set this environment variable to the address of the IFRT Proxy server when using the proxy backend.
Set in jax.distributed.initialize for oneAPI GPU support.
Set to proxy to use the newly registered backend that connects to an IFRT Proxy server.
Implemented reverse iteration using chunked slicing, matching __iter__ for faster eager-mode reverse iteration.
Added a remat option for writing jax.custom_vjp.defremat rules with closures.
Customize how a custom VJP function is rematerialized when differentiated under jax.remat; it requires the jax_custom_vjp3 and jax_remat3 implementations.
Added an option for batching under jax.vmap when a primal output is shared across the batch.
Added a way to query the default device layout for an array.
Added an argument to place optimization barriers on the first input and/or last output of the scope.
Added tools to author, inspect, splice, and execute XLA:GPU ThunkExecutor programs directly from Python.
Added support for accuracy=jax.lax.AccuracyMode.HIGHEST, with improved worst-case accuracy and accurate first- and higher-order derivatives.
Added polynomial evaluation with a primitive that is faster and uses less memory when computing gradients.
Accepts mxfp4, using E2M1 elements with one E8M0 scale per block of 32; quantization runs natively on AMD gfx950 GPUs and uses a dequantize-and-dot fallback elsewhere.
Added batch_size, defaulting to 16, to control the batched approach now used by default for many-key sorts and avoid long compile times.
Added profile data export for session- and context-manager-based profiling.
Added session-based profiling.
Introduced as an Array type alias; starting in JAX 0.13.0, it will be the return type for key, split, and fold_in.
Added the inverse complementary error function.
Added a context manager that raises an error at trace time in explicit sharding mode if an operation or reshard introduces a collective across a specified mesh axis.
Performance 3
Now lowers via a single-pass jax.lax.reduce.
Uses an unstable sort when weights=None, avoiding a companion index sort.
Uses an unstable sort when weights=None, avoiding a companion index sort.
Behavior and default changes 1
Sets jax_oneapi_visible_devices in oneAPI GPU support.
Tooling 2
Now annotated as returning Iterator[Self] rather than Any.
Now annotated as returning Iterator[Self] rather than Any.
Fixes and other changes 23
Fixed a regression where cached host-transfer CPU buffers prevented buffer donation.
Fixed indexing and iteration over fully replicated multi-device arrays.
A failing ordered io_callback no longer permanently poisons effect tokens.
Complex absolute value has improved numerical accuracy by lowering through power-of-two scaled hypot.
Improved precision of fallback expansions near zero under flush-to-zero mode.
Fixed numerical overflow in gradients for large inputs.
Improved CPU accuracy for float32 and float64 and tightened float32 accuracy on specified TPU generations.
Tightened float32 accuracy on TPU v5p, v6e, and 7x with libtpu >= 0.0.50.
Fixed incorrect results for signed integer minimum values such as INT_MIN.
Preserves FP8/FP4 quantization when XLA expands scaled matmul into dequantize-and-dot and computes MX block scales in float32 regardless of input dtype.
Complex absolute value has improved numerical accuracy by lowering through power-of-two scaled hypot.
Fixed incorrect results for signed integer minimum values such as INT_MIN.
Fixed the gradient at x = 0.0, which previously returned 1.0 instead of 2**n.
Fixed incorrect results for signed integer minimum values such as INT_MIN.
Improved numerical accuracy across floating-point types.
Now clears the caller's concrete and abstract mesh when invoking host callbacks, matching jax.debug.callback.
Fixed parameter broadcasting.
Now accepts batched inputs of shape (..., M, N), as documented.
Fixed the Hessian when a == b.
Reimplemented erfcx, improving accuracy and derivatives as part of the special-function accuracy changes.
Numerical accuracy of ndtri and its derivatives improved using minimax polynomial approximations and a custom JVP rule.
Fixed vmap of jax.lax.scan closing over a Ref batched on a non-leading axis.
Fixed out-of-bounds slice validation when start = 0.
jax 0.1166 changes
Breaking changes 13
JAX Bazel builds now use Bzlmod instead of WORKSPACE.
The minimum CuDNN version required for CUDA 12 is now v9.10.2.
Removed the deprecated API; use explicit sharding mode instead.
inline=True now corresponds to jax.Inline.JAX_LATE instead of jax.Inline.JAX_EARLY.
An input empty along the reduction axis now raises ValueError instead of an internal gather error.
When dtype is unspecified, the returned array now uses the default float dtype rather than always float32.
The flag has been removed in favor of the
EffortLevelenum.Deserialization of Exported modules from before January 15th, 2026 is no longer supported; NamedSharding is now the only supported sharding serialization.
Now returns a tuple rather than a list, aligning with NumPy>2.0 and the Array API specification.
Now returns a tuple rather than a list, aligning with NumPy>2.0 and the Array API specification.
Now returns a tuple rather than a list, aligning with NumPy>2.0 and the Array API specification.
The default is now always
True;Noneis no longer a valid value.The flag has been removed in favor of the
EffortLevelenum.
Security 1
Added mutual TLS coordination-service configuration arguments and corresponding environment variables.
Deprecations 2
Accessing this field now raises a warning; use
in_shardings_jaxinstead.Accessing this field now raises a warning; use
out_shardings_jaxinstead.
New APIs and options 10
jaxlib wheels now include C++ FFI extension headers to support out-of-tree plugins.
Added the exclude_argnames argument.
Added an API for querying conservative bounds on symbolic dimension expressions.
Added log2 as a first-class primitive; jax.numpy.log2 now lowers through it.
Added the out_sharding parameter.
Added an API that returns both the minimum and maximum of an array, matching NumPy 2.3+.
Added oneMKL LAPACK kernels for QR, LU, symmetric/Hermitian eigenvalue, SVD, and hybrid solvers on oneAPI GPUs.
JAX now reads this environment variable in GKE TPU clusters.
Temporarily bypasses the error check when deserializing JAX exports older than the backwards compatibility window.
Adds
jax.numpy.top_k, implementingnumpy.top_k.
Performance 3
Scan-based lowering was re-enabled on GPU, improving performance.
GPU multidimensional inverse real FFTs again lower to a single C2R transform; results are unchanged and typical sizes are about 1.4x faster.
Large Mesh construction avoids redundant device array allocations and copies.
Behavior and default changes 3
Cluster detection now supports Open MPI 5.
The p parameter now accepts RealArray rather than only float, allowing array-valued shape parameters.
cuDNN fused attention batching rules now support operands without the vmap axis, including shared bias or
mask.
Tooling 1
JAX now uses Bazel 8.7.0 to build from source.
Fixes and other changes 33
Improved thread safety across JAX and jaxlib for free-threaded Python 3.13t, 3.14t, and 3.15t.
JAX pytrees now support frozendict on Python 3.15.
Fixed tangents when auxiliary values are integer-typed.
Fixed color cycling.
Error code assignment is now deterministic.
Fixed symbolic product bounds when even-powered factor bounds cross zero or zero factors are paired with infinite bounds.
Fixed first-, second-, and higher-order autodiff gradients at 0.0.
Adjusted the JVP rule so it does not depend on bitwise equivalence between forward- and backward-pass results.
Fixed a potential division-by-zero error for zero-size inputs.
Fixed numerical overflow in gradients for large inputs.
Fixed a crash on empty arrays.
With size=0, now returns an empty array of the natural result dtype instead of raising ValueError.
For singular matrices with p=None or 2, now returns infinity instead of NaN, matching NumPy and the other norms.
Fixed gradients that produced NaN or incorrect values for large eigenvalues.
Numerical accuracy improved by pre-computing the 1 / log(2) constant factor.
With size=0 on non-empty inputs, now returns an empty array instead of raising IndexError.
Now uses a Taylor series near zero, giving accurate derivatives of all orders.
With axis specified, empty-axis arrays now have NumPy's output shape instead of a fabricated phantom slice for fully empty inputs.
The argument signature now accepts bool | Sequence[bool], matching jax.checkpoint.
Fixed gradients leaking into the unused triangle when symmetrize_input=False.
Fixed state discharge accumulation on indexed ReshapeTransform views.
Worked around a rocFFT twiddle cache bug in multidimensional real FFTs.
Setting this to
Falseaddsrematerializationto disabled XLA passes instead of overwriting the set, preserving passes disabled throughXLA_FLAGS.Errors for non-scalar function outputs now suggest reducing the output to a scalar, using
jax.jacobian, or reshaping size-1 outputs.Errors for non-static or traced slice indices now suggest
jax.lax.dynamic_slice,jax.lax.dynamic_update_slice, orjax.ds, and show tracer provenance.Fixes propagation of singleton sharded dimensions when reshaping arrays with sharding constraints.
Abstract evaluation now checks
.matequivalency only when the abstract value is aShapedArray.The cuDNN fused attention backward pass no longer computes a bias gradient when the only attention bias comes from a boolean
mask; explicit bias and non-boolean mask gradients are unchanged.Uses closed-form LU decomposition with row pivoting for 2x2 and 3x3 matrices instead of polynomial expansions to avoid numerical instability and catastrophic cancellation.
Uses closed-form LU decomposition with row pivoting for 2x2 and 3x3 matrices instead of polynomial expansions to avoid numerical instability and catastrophic cancellation.
Fixed for
namedtupleinstances.Errors for non-scalar function outputs now suggest reducing the output to a scalar, using
jax.jacobian, or reshaping size-1 outputs.Vmapping fp8 cuDNN fused attention now works; amax outputs are whole-batch statistics without the vmap axis, while vmapping scale/descale operands raises
NotImplementedError.