When a function has several inputs and outputs, one derivative is not enough. A Jacobian collects every first-order partial derivative of that vector-valued function into a matrix. section 4.2.5
JAX provides jacfwd() and jacrev(). Both return the same Jacobian values, but forward-mode autodiff suits tall matrices with many outputs, while reverse-mode autodiff suits wide matrices with many inputs. section 4.2.5 section 4.3.2
Computing a vector-valued Jacobian
For f(x) = [x₀² + x₁² − x₁x₂, x₀² − x₁² + 3x₀x₂] at x = [3, 4, 5], both transformations produce [[6, 3, −4], [21, −8, 9]]. The example reports jacfwd() faster for this two-output, three-input function. Use jvp() instead when you need a Jacobian-vector product without materializing the full matrix. section 4.2.5 section 4.3.2
The Jacobian connects several levels of automatic differentiation:
- Forward mode computes one Jacobian column per input direction and can recover the full matrix in as many evaluations as inputs. section 4.3.2
- The chain rule expresses derivatives through products of layer Jacobians, connecting Jacobians to neural-network differentiation. section 4.2.5
