jax-ml/jax jax-v0.12.0

JAX v0.12.0

on GitHub

New features

• Added jax.scipy.special.erfcinv for the inverse complementary
  error function.
• Added jax.lax.polynomial for polynomial evaluation. The polynomial
  primitive is faster and uses less memory when computing gradients.
• Added jax.custom_vjp.defremat, which customizes how a
  jax.custom_vjp function is rematerialized when it's differentiated
  under jax.remat. It requires the jax_custom_vjp3 and jax_remat3
  implementations.
• Added a remat option to jax.custom_gradient, for writing
  jax.custom_vjp.defremat rules with closures.
• arr.at[...].get() now accepts a strategy argument that allows choosing
  whether the operation lowers to gather, dynamic_slice, or slice.
• Added a generic_batching option to jax.custom_vjp.defvjp, for
  batching under jax.vmap when a primal output is shared across the
  batch. See jax-301-vmap-shared-outputs.
• jax.nn.get_scaled_dot_general_config accepts 'mxfp4': E2M1
  elements with one E8M0 scale per block of 32, for
  jax.nn.scaled_dot_general. The quantization is executed natively
  on AMD gfx950 GPUs and by a dequantize-and-dot fallback elsewhere.
• JAX now registers a proxy backend, which connects to an IFRT Proxy server
  such as Pathways on Cloud.
  To use it, set JAX_PLATFORMS=proxy and set JAX_BACKEND_TARGET to the
  server address, e.g. grpc://localhost:29000.
• Added 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 any specified
  mesh axis.
• Added jax.custom_vjp.defremat_with_logs, allowing
  jax.custom_vjp.defremat rules (and jax.custom_gradient with
  remat=True and with_logs=True) to log values from the backward pass
  (#40956).
• jax.scipy.linalg.eigh now supports generalized Hermitian
  eigenproblems (b != None) for problem types type=1, 2, and 3
  (#41162, #41197).
• Added jax.profiler.ProfileSession and
  jax.profiler.export_profile_data for session- and
  context-manager-based profiling.
• Added jax.experimental.layout.Layout.for_array to query the default
  device layout for an array.
• Added a strict_in_out argument to jax.experimental.program_order
  to place optimization barriers on the first input and/or last output of the
  scope.
• Added jax.experimental.thunky for authoring, inspecting, splicing,
  and executing XLA:GPU ThunkExecutor programs directly from Python.
• Added JVP rule for jax.lax.nextafter, enabling forward- and
  reverse-mode autodiff through nextafter (#38679).
• Widened jax.random.ball's p parameter type from float to
  RealArray, allowing array-valued norm orders (#41546).
• Added support for DLPack 1.0 (DLManagedTensorVersioned) and read-only
  DLPack tensors (#40789).
• [oneAPI GPU] Added SYCL/oneMKL kernels for
  jax.lax.linalg.lu_pivots_to_permutation,
  jax.lax.linalg.cholesky_update,
  jax.lax.linalg.tridiagonal_solve (with perturb_singular=True),
  jax.lax.linalg.tridiagonal (sytrd/hetrd), symmetric rank-$k$
  updates (syrk), and scaled_matmul, and set jax_oneapi_visible_devices
  in jax.distributed.initialize.
• [ROCm] Added gfx1250 to the default target architectures in ROCm builds
  (#40877).

Breaking changes

• Removed jax.custom_remat. Use jax.custom_vjp.defremat instead:
  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
  (roughly 2013 onwards, requiring AVX2 and FMA support). Older CPUs that
  only support AVX are no longer supported, though you can still build from
  source for older CPUs.
• jax.numpy.atleast_1d, jax.numpy.atleast_2d, and
  jax.numpy.atleast_3d now return tuples of arrays rather than
  lists, matching the behavior in NumPy 2.0+.
• The shape argument of jax.random.multinomial is now the batch
  shape, as documented and as in NumPy: the result has shape
  (*shape, p.shape[-1]). Previously shape had to include the outcome
  axis. Code that passed the full output shape should drop its last entry.
• Under jax.vmap, a jax.custom_vjp function with a primal
  output that doesn't depend on any batched input could get gradients that
  counted that output's cotangent once per example. That case now raises an
  error explaining how to resolve it, for example with the new
  generic_batching option. See jax-301-vmap-shared-outputs.
• Dropped support for NumPy 2.1; JAX now requires NumPy 2.2 or newer.
• jax.numpy.cross no longer accepts 2D vector inputs, completing a
  deprecation started in JAX 0.11.0 (#40876).
• The copy, order, and ndmin arguments to jax.numpy.array are
  now keyword-only, completing a deprecation started in an earlier release
  (#40878).
• jax.profiler.ProfileOptions now raises AttributeError when
  unknown attributes are assigned to it.
• jax.checkpoint, jax.custom_vjp, and jax.shard_map
  now raise an error if the same Ref is passed as more than one argument or
  is both closed over and passed as an argument.

Deprecations

• jax.lax.psend and jax.lax.precv are deprecated.

Changes

• Improved the numerical accuracy of jax.scipy.special.ndtri and its
  derivatives by replacing the piecewise rational approximations with
  minimax polynomial approximations and a custom JVP rule.
• JAX now uses Bazel 9.2.0 to build from source.
• Improved the numerical accuracy of complex jax.numpy.abs and
  jax.lax.abs by lowering through power-of-two scaled hypot.
• Improved the numerical accuracy of jax.numpy.hypot and its
  derivatives by replacing the division-based formula with power-of-two
  scaling and a custom JVP rule, avoiding rounding errors for large inputs.
  The derivative at (0, 0) now returns NaN (previously ±0.5).
• Improved the numerical accuracy of jax.numpy.sinc across all
  floating-point types (to $\le 0.5$ ULP for bfloat16 and float16,
  $\le 2.5\text{–}4.0$ ULPs for float32, and $\ge 2.0$ ULPs for float64).
• jax.numpy.ldexp and jax.numpy.frexp are now bit-exact for
  normal floating-point numbers, replacing floating-point arithmetic with
  bitwise operations. Subnormal inputs and underflowing outputs are flushed to
  signed zero.
• Several jax.numpy functions are now implemented in terms of HiJax
  primitives (see jax-301-hijax-types). This includes
  ~jax.numpy.searchsorted, ~jax.numpy.nonzero,
  ~jax.numpy.ldexp, jax.numpy.frexp. Typical users should not
  see any change in behavior, but you may notice more efficient batching
  and autodiff, and more concise representations in jaxprs.
• The "web" debugger now defaults to hostname "localhost", which is a safer
  default that avoids inadvertently opening a port to the world.
• jax.numpy.lexsort with many keys now uses a batched approach by
  default, controlled by the new batch_size parameter (which defaults to
  16). This avoids long compile times for sorts with many keys. Fixes the
  poor performance of jax.numpy.unique for high-dimensional arrays
  (#17370).
• Reimplemented jax.scipy.special.erfcx and added
  accuracy=jax.lax.AccuracyMode.HIGHEST support for jax.lax.erfc,
  improving worst-case float32 accuracy to $\le 3.5\text{–}5.5$ ULPs across
  CPU, GPU, and TPU (and $\le 2.5\text{–}3.0$ ULPs in float64) with accurate
  first- and higher-order derivatives.
• Improved float32 accuracy of jax.lax.lgamma
  (jax.scipy.special.gammaln) to $\le 4$ ULPs across CPU, GPU, and TPU
  (#21284).
• Improved float32 and float64 accuracy of jax.lax.expm1 on CPU to
  $< 1$ ULP, and tightened expm1 and log1p float32 accuracy on TPU v5p,
  v6e, and 7x (with libtpu >= 0.0.50).
• Improved precision of jax.lax.asin fallback expansions near zero
  under flush-to-zero mode.
• jax.scipy.special.betainc now stops its continued fraction iteration
  per element as each element converges rather than running until all batch
  elements converge (#28547).
• jax.pure_callback and jax.experimental.io_callback now clear
  the caller's concrete and abstract mesh when invoking host callbacks,
  matching jax.debug.callback.
• jax.numpy.quantile and jax.numpy.percentile now use an
  unstable sort when weights=None, avoiding a companion index sort.
• jax.numpy.minmax now lowers via a single-pass jax.lax.reduce.
• Implemented __reversed__ on jax.Array using chunked slicing,
  matching __iter__ for faster reverse iteration in eager mode.
• Updated minimum pinned libtpu version to 0.0.50.
• When cache key equality comparison fails in weakref_lru_cache, JAX now
  raises a descriptive ValueError reporting the non-comparable key types
  (#21825).

Type annotation improvements

• jax.numpy.int32, jax.numpy.float32, and other scalar constructors
  now have proper type annotations; prior to v0.12.0 they were typed Any.
  This may lead to new typecheck errors in cases where they were being misused;
  for example, if they were assumed to be instances of jnp.dtype (which they
  are not) or used as annotations for arrays of a given type (which is invalid).
• jax.Array.__iter__ and jax.Array.__reversed__ are now annotated
  as returning Iterator[Self] rather than Any.
• jax.random.Key was introduced as an Array type alias. Starting in JAX
  v0.13.0, this will be the return type for jax.random.key,
  jax.random.split, and jax.random.fold_in.
• Unary jax.Array methods, jax.Array.astype, and
  arr.at[...].set() now use Self return type annotations
  (#41055, #41324).

Bug fixes

• Fixed the gradient of jax.numpy.ldexp at x = 0.0, which previously
  returned 1.0 instead of 2**n.
• jax.nn.scaled_dot_general keeps its FP8/FP4 quantization when XLA
  expands the scaled matmul into a dequantize-and-dot (previously the
  rounding could be folded away under xla_allow_excess_precision) and
  computes MX block scales in float32 regardless of the input dtype.
• Fixed the derivative of jax.lax.sign and jax.numpy.sign for
  complex inputs, which was previously always zero (#41000).
• jax.scipy.linalg.lu now accepts batched inputs of shape
  (..., M, N), as documented; previously it raised a ValueError for any
  input with more than two dimensions.
• Fixed reverse-mode autodiff for jax.experimental.sparse.coo_matvec,
  jax.experimental.sparse.coo_matmat,
  jax.experimental.sparse.csr_matvec, and
  jax.experimental.sparse.csr_matmat when transpose=True
  (#40625).
• Fixed jax.numpy.linalg.svd with hermitian=True breaking unitarity
  when a singular value is zero (#41547).
• Fixed gradient cancellation in jax.lax.expm1, jax.nn.elu,
  and jax.nn.selu for large negative inputs (#39794,
  #39795, #39796).
• Fixed numerical overflow in jax.lax.asinh gradients for large
  inputs.
• Fixed precision, overflow, and NaN/inf handling in
  jax.scipy.special.rel_entr and jax.scipy.special.kl_div
  (#38239, #38262, #41200).
• Fixed domain, boundary, and overflow issues in
  jax.scipy.special.expn and jax.scipy.special.expi at x = 0
  (#41221).
• Fixed jax.scipy.special.betaln Hessian when a == b.
• Fixed jax.scipy.special.zeta's JVP rule returning finite primal
  values outside its domain (x < 1) where zeta returns NaN
  (#40968).
• Fixed jax.scipy.stats.norm.isf losing precision and returning inf
  for small tail probabilities (#41262).
• Fixed jax.scipy.stats.expon.logcdf returning -inf for small
  positive inputs (#40965).
• Fixed NaN gradients in jax.scipy.stats.laplace.cdf in the tails
  (#40966) and in jax.scipy.stats.poisson.entropy
  (#40969).
• Fixed domain handling and tail accuracy in jax.scipy.stats.cauchy.ppf
  and jax.scipy.stats.cauchy.isf (#38318).
• Fixed jax.scipy.stats.vonmises.logpdf at kappa = 0
  (#38621) and jax.scipy.stats.chi2.logpdf at x = 0 and
  x = inf for df = 2 (#38316).
• Fixed jax.scipy.stats.multinomial.logpmf and pmf for batched x
  and n (#38605), jax.scipy.stats.poisson.logpmf for
  fractional loc (#41220), jax.scipy.stats.pareto
  support checks when scale is below the floating-point resolution of loc
  (#40964), and NaN propagation in
  jax.scipy.stats.uniform and jax.scipy.stats.truncnorm
  (#38317, #38606).
• Fixed parameter broadcasting in jax.random.f.
• Fixed jax.Array.sum to forward the initial parameter to
  jax.numpy.sum (#41137).
• Fixed jax.numpy.ufunc.reduceat ignoring the specified dtype when
  casting inputs before accumulation (#41247).
• Fixed jax.numpy.partition, jax.numpy.argpartition, and
  jax.lax.top_k returning incorrect results for signed integer
  minimum values (e.g. INT_MIN).
• Fixed transpose rule for jax.lax.select_n with out-of-range integer
  indices or a single case (#10993).
• Propagated mode, indices_are_sorted, and unique_indices through
  jax.lax.scatter_max and jax.lax.scatter_min JVP rules
  (#40626).
• Fixed input shape validation in jax.numpy.linalg.eigh,
  jax.scipy.linalg.eigh, and jax.numpy.linalg.cholesky, and
  raised NotImplementedError for unimplemented JVP of
  jax.lax.linalg.svd with partial subset_by_index and
  compute_uv=True (#41210).
• Fixed handling of negative static_argnums and donate_argnums with named
  parameters in jax.jit (#40988), and fixed keyword
  argument forwarding in jax.fwd_and_bwd (#41156).
• Fixed indexing and iteration over fully replicated multi-device arrays and
  out-of-bounds slice validation when start = 0.
• Fixed jax.vmap of jax.lax.scan closing over a Ref batched
  on a non-leading axis.
• Fixed jax.experimental.checkify error code dtype mismatches under
  nested jax.enable_x64() scopes (#40482).
• Fixed a regression where cached host-transfer CPU buffers prevented buffer
  donation, and prevented a failing ordered io_callback from permanently
  poisoning effect tokens.