Newly documented

Contents

Newly documented#

Features documented for the first time. See also the Change log.

August 2026#

  • Autodiff and sharding. How cotangent shardings follow from primal shardings in explicit mode, and controlling backward-pass communication with unreduced and reduced shardings.

  • Refs: mutable arrays. jax.new_ref creates an array ref that can be read and written in place, composing with transformations. In-place updates under jit in In-place updates with refs; autodiff with refs in Autodiff with refs.

  • First-class VJP objects. The callable returned by jax.vjp() is a pytree. A recipe for getting the forward and backward passes as separate functions.

  • saveable_args on jax.vjp. Exclude argument values (like weights) from what a VJP saves.

  • Custom derivatives with hijax primitives. One primitive can carry rules for both differentiation modes and batching. More capable alternative to jax.custom_vjp and jax.custom_jvp.

  • Backward-pass logging. Plumb data out of backward passes, e.g. for gradient diagnostics.

  • Structured residuals. Organize what the forward pass saves for the backward pass.

  • New JAX types with hijax. Define new types with their own tangent types, batching behaviors, and sharding. Consume with your own hijax primitives.

  • FFI with hijax. The foreign function interface docs are rewritten around hijax primitives, so foreign calls can carry their own rules for batching, differentiation, and sharding — composing with vmap, grad, and partitioned inputs.

  • Fault tolerance. Surviving machine failures in multi-host jobs with jax.live_devices.

  • Compiler control. Compilation effort levels per jit-compiled function, and per-operation XLA metadata.

  • Matrix multiplication precision control. How to control matrix multiplication algorithms, on each operation and globally.

Contents