Automatic differentiation computes derivatives of code-expressed computations by tracking elementary operations and propagating derivatives with the chain rule. It handles complex programs, including branching, loops, recursion, and nested neural-network layers, without requiring manually derived formulas. section 4.1.4 section 4.4
In JAX, automatic differentiation is a composable transformation of Python and NumPy-style functions. Calling grad() produces a gradient function, while related transformations provide function values with gradients, Jacobians, Hessians, and higher-order derivatives. section 1.2.2 section 4.4
Choose an autodiff mode according to the function's input-output shape:
- Forward mode computes derivatives of all outputs with respect to one input in a single pass and suits functions with many outputs. section 4.4
- Reverse mode computes the gradient of one output with respect to all inputs in a single pass and suits functions with many inputs and few outputs. section 4.4
- Use jvp() for a Jacobian-vector product or vjp() for a vector-Jacobian product without explicitly constructing the full Jacobian. section 4.3 section 4.4
Differentiate a Python function with JAX
For f(x) = x**4 + 12*x + 1/x, define df = jax.grad(f) and evaluate df(x). At x = 11.0, JAX returns approximately 5335.9917, avoiding a hand-written derivative function. The resulting gradient can support gradient descent, which updates neural-network weights opposite the gradient to reduce loss. section 4.1 section 4.1.4 section 4.0
Figure 1.3 JAX (right) is much more than NumPy (left). It provides a set of composable function transformations for Python + NumPy code for compilation, vectorization, parallelization, and automatic differentiation.
Use the transformation map to see automatic differentiation alongside JAX compilation, vectorization, and parallelization, and to understand why these operations can be composed. section 1.2.2
