New features Added a doc on defining custom derivative rules with the experimental hijax API (hijax-custom-derivatives), along with jax.experimental.hijax helpers for deriving VJPHiPrimitive autodiff rules from a jvp or lin rule: linearize_from_jvp with apply_derived_linearization, vjp_fwd_from_jvp with transpose_jvp, vjp_fwd_from_lin with transpose_linearized, and jvp_from_lin. Added jax.custom_r…
New features Added jax.scipy.linalg.invhilbert for the closed-form inverse of the Hilbert matrix ({jax-issue}#10144). Added jax.scipy.linalg.invpascal for the inverse of the Pascal matrix (#10144). Added jax.scipy.linalg.fiedler_companion for constructing the pentadiagonal Fiedler companion matrix of a polynomial (#10144). Added jax.ShapeDtypeStruct.like -- a shortcut for constructing a jax.ShapeD…
New features Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098). Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144). Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144). Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144). Added jax.scipy.linalg.leslie for constructin…
New features: Added ResizeMethod.CUBIC_PYTORCH to jax.image.resize to match PyTorch's bicubic resize (#15768). We now support differentiation of jax.lax.linalg.qr for wide matrices and when full_matrices is True. LAPACK operations are now parallelized along the batch dimension on CPU. Added perturb_singular argument to jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near…
JAX 0.9.2 (March 2, 2026) Changes: The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it. jax.numpy.arange with step specified no longer generates the array on host. The benefit is more efficient code, though this can lead to less precise outputs for narrow-width floats (e.g. bfloat16). To recover the previous behavior in this case, use…
Changes: JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array. Using jax.shard_map in Explicit mode will raise an error if the PartitionSpec of input does not match the PartitionSpec specified in in_specs. In other words, it will act like an assert instead of an implicit reshard. in_specs is an optional argument so you can omit specif…
JAX v0.9.0.1 is identical to v0.9.0 with the commits from the following four PRs patched in: openxla/xla#36579 openxla/xla#36345 openxla/xla#36755 openxla/xla#36696
JAX v0.8.3 is identical to v0.8.2 with the following two bug fixes patched in: openxla/xla@4bc723d openxla/xla@21552fd
New features: Added jax.thread_guard, a context manager that detects when devices are used by multiple threads in multi-controller JAX. Bug fixes: Fixed a workspace size calculation error for pivoted QR (magma_zgeqp3_gpu) in MAGMA 2.9.0 when using use_magma=True and pivoting=True. (#34145). Deprecations: The flag jax_collectives_common_channel_id was removed. The jax_pmap_no_rank_reduction config…
Deprecations jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement. Complex arguments passed to jax.numpy.arange now result in a deprecation warning, because the output is poorly-defined. From jax.core a number of symbols are newly deprecated including: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_par…
New features: jax.jit now supports the decorator factory pattern; i.e instead of writing @functools.partial(jax.jit, static_argnames=['n']) def f(x, n): ... you may write @jax.jit(static_argnames=['n']) def f(x, n): ... Changes: jax.lax.linalg.eigh now accepts an implementation argument to select between QR (CPU/GPU), Jacobi (GPU/TPU), and QDWH (TPU) implementations. The EighImplementation enum is…