New features
- Added jax.numpy.minmax (and jnp.minmax), which returns both the minimum and maximum of an array, matching NumPy 2.3+ (#40089).
- Added jax.lax.log2 and primitive jax.lax.log2_p, making log2 a first-class primitive in JAX (jax.numpy.log2 now lowers via jax.lax.log2).
- Added jax.lax.one_minus_square primitive to accurately compute 1 - x^2 near $\pm 1$ and provide accurate derivatives near $0$.
- Added jax.export.symbolic_dim_bounds for querying conservative bounds on symbolic dimension expressions (#40006).
- Added frozendict support to JAX pytrees for Python 3.15 (PEP 814).
- jax.distributed.initialize can now secure the coordination service with mutual TLS via the new mtls_cert_file, mtls_key_file, mtls_ca_file, mtls_peer_uri_prefix and verify_secure_credentials arguments (or the JAX_MTLS_CERT_FILE, JAX_MTLS_KEY_FILE, JAX_MTLS_CA_FILE, JAX_MTLS_PEER_URI_PREFIX and JAX_DISTRIBUTED_VERIFY_SECURE_CREDENTIALS environment variables).
- Added cluster detection support for Open MPI 5 in jax.distributed.initialize (#40512).
- Added support for reading TPU_PROCESS_ADDRESSES_PATH in GKE TPU clusters.
- Widened jax.random.generalized_normal's p parameter type from float to RealArray, allowing array-valued shape parameters (#40126).
- Added exclude_argnames argument to jax.experimental.program_order.
- [oneAPI GPU] Added oneMKL LAPACK solver kernels for QR decomposition (geqrf, orgqr/ungqr, ormqr/unmqr), LU decomposition (getrf), symmetric/Hermitian eigenvalue decomposition (syevd/heevd), SVD (gesvd), and hybrid solver kernels (geqp3, eig) (#40000, #40186, #40543).
- jaxlib wheels now ship C++ FFI extension headers (collectives.h, record.h) to support out-of-tree plugins (#40333).
Breaking changes
- Removed deprecated jax.experimental.shard_alike. Use explicit sharding mode instead (see sharding).
Changes
- Optimized large JAX jax.sharding.Mesh construction by avoiding redundant device array allocations and copies.
- Thread-safety improvements across JAX and jaxlib for free-threaded Python (Python 3.13t, 3.14t, 3.15t).
- inline=True in jax.jit now corresponds to jax.Inline.JAX_LATE instead of jax.Inline.JAX_EARLY.
- The minimum CuDNN version for CUDA 12 is v9.10.2.
- JAX now uses Bazel 8.7.0 to build from source.
- JAX now uses Bzlmod for its Bazel builds instead of WORKSPACE.
- On GPU, multi-dimensional inverse real FFTs (jax.numpy.fft.irfftn, jax.numpy.fft.irfft2 and jax.lax.fft with FftType.IRFFT) are again lowered to a single C2R transform, as before JAX 0.10.0, instead of an IFFT over the outer axes and a 1-D IRFFT with two transposes. The input is first made Hermitian-symmetric along the outer axes, which does not change the result under NumPy's convention (only the last axis is assumed symmetric), so results are unchanged while the transform is ~1.4x faster at typical sizes.
- Re-enabled scan-based lowering of cumulative sums (jax.numpy.cumsum) on GPU, improving performance.
- jax.numpy.tri now returns an array with the default float dtype when the dtype argument is not specified. Previously it always returned float32 (#40242).
- jax.numpy.unique with axis specified now matches NumPy's output shape for arrays that are empty along the given axis, instead of fabricating a phantom slice for fully-empty inputs.
- Improved numerical accuracy of jax.numpy.log2 by pre-computing the 1 / log(2) constant factor (#40430).
- Added out_sharding parameter to jax.numpy.histogram.
- Updated jax.remat's prevent_cse argument signature to accept bool | Sequence[bool], matching jax.checkpoint.
- Made jax.experimental.checkify error code assignment deterministic.
Bug fixes
- Fixed numerical overflow in jax.numpy.arccosh and jax.lax.acosh gradients for large inputs (#40643, #40634).
- Fixed first-, second-, and higher-order autodiff gradients of jax.lax.bessel_i0e and jax.lax.bessel_i1e at 0.0 (#40640, #40635).
- Fixed jax.numpy.linalg.eigh gradients producing NaN or incorrect values for large eigenvalues (#40149, #40141).
- jax.numpy.sinc now uses a Taylor series near zero, giving accurate derivatives of all orders. Previously, autodiff of the sin(πx)/(πx) quotient suffered catastrophic cancellation near zero (#34139, #10750).
- Fixed a bug where jax.numpy.linalg.cond returned NaN instead of infinity for singular matrices when p is None or 2, matching NumPy and the other norms.
- Fixed jax.numpy.histogram crashing on empty arrays (#40025, #40020).
- Fixed jax.numpy.intersect1d and jax.numpy.setxor1d with size=0, which previously raised a ValueError; they now return empty arrays of the natural result dtype.
- Fixed jax.numpy.setdiff1d raising an IndexError when called with size=0 on non-empty inputs; it now returns an empty array.
- Fixed incorrect gradients for jax.scipy.linalg.cholesky and jax.numpy.linalg.cholesky with symmetrize_input=False where non-zero gradients leaked into the unused triangle of the input matrix (#40421).
- Fixed jax.numpy.median on an input that is empty along the reduction axis, which previously raised an internal error from gather; it now raises a ValueError.
- Fixed a potential division-by-zero error in jax.nn.initializers.variance_scaling for zero-size inputs (#35096).
- Fixed jax.custom_root tangents when auxiliary values are integer-typed (#39913, #24295).
- Fixed color cycling in jax.debug.visualize_array_sharding (#39922, #25695).
- Adjusted JVP rule for jax.lax.min and jax.lax.max to not depend on bitwise equivalence between forward and backward pass results (#40578).
- Fixed symbolic product bounds calculation in jax.export when even-powered factor bounds cross zero or zero factors are paired with infinite bounds (#40054).
- Fixed state discharge accumulation (ref_addupdate) on indexed ReshapeTransform views.
- [ROCm] Worked around a rocFFT twiddle cache bug in multi-dimensional real FFTs (#40389).