JAX 301: advanced autodiff and extending JAX

The 101 docs introduced jax.grad(); the pages here go as deep as you like: the machinery underneath (JVPs, VJPs, Jacobians, Hessians), how autodiff interacts with sharding, defining your own derivative rules and even your own types, autodiff with mutable state, and controlling the memory/compute tradeoff of differentiation.

1. The Autodiff Cookbook with JVP and VJP — autodiff’s basic ingredients, the jvp and vjp machinery: recipes for Hessian-vector products and full Jacobians, plus vmap and complex numbers.

2. First-class VJPs — the VJP object as a pytree, splitting the forward and backward passes into separately compiled functions run on your own schedule, and excluding argument values (like weights) from the saved state with saveable_args.

3. Autodiff and sharding — how autodiff interacts with sharding: cotangent shardings as a function of primal shardings, and controlling backward-pass communication with unreduced and reduced, in both explicit and manual (shard_map) modes.

4. Custom derivative rules with hijax primitives — defining custom derivative rules with hijax primitives, the recommended approach: one primitive can carry rules for both modes, plus linearization, batching, producing structured residuals, logging data out of the backward pass, and more.

5. Custom JVPs or VJPs with custom_jvp and custom_vjp — the classic decorators, customizing one differentiation mode at a time: still fully supported, and often the most convenient tool for simple cases.

6. Autodiff with refs — autodiff with mutable arrays: plumbing values out of backward passes, in-place gradient accumulation with with_refs, and differentiating with respect to refs.

7. Gradient checkpointing with jax.checkpoint (jax.remat) — gradient checkpointing with jax.checkpoint: what autodiff saves versus recomputes, name-based policies, offloading, and per-function control with custom_remat.

8. Defining new JAX types with hijax — defining entirely new JAX types with hijax, with their own derivatives, batching, and sharding behavior.