Automatic differentiation is a core component of JAX, enabling gradient-based optimization for machine learning models. While jax.grad provides a convenient interface for obtaining gradients, understanding and controlling the underlying mechanisms offers greater flexibility and efficiency for complex tasks.

This chapter moves beyond the basics of jax.grad. You will learn about the two fundamental modes of automatic differentiation: forward-mode (Jacobian-vector products, JVPs) and reverse-mode (vector-Jacobian products, VJPs). We will cover how to compute these directly using jax.jvp and jax.vjp, which form the building blocks for jax.grad and other transformations.

Key topics include:

• Calculating higher-order derivatives by composing differentiation transformations.
• Strategies for efficiently computing full Jacobian and Hessian matrices.
• Defining custom differentiation rules for functions using jax.custom_vjp and jax.custom_jvp, useful for numerical stability, performance optimization, or handling non-JAX code.
• Understanding how automatic differentiation interacts with control flow primitives like lax.scan, lax.cond, and lax.while_loop.
• Techniques for handling functions that are not differentiable or where gradients should be explicitly stopped, such as using jax.lax.stop_gradient.

By the end of this chapter, you will have a more detailed understanding of JAX's autodiff system and the tools to apply it effectively in advanced scenarios.