• 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).