Skip to content

JAX v0.11.2

Choose a tag to compare

View all tags
@hawkinsp hawkinsp released this 17 Sep 23:43
jax-v0.11.2
3254480
  • 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).
Assets 2
Loading