When forward-mode autodiff evaluates a function, it carries each intermediate value as a pair: a primal value and a tangent. The tangent records sensitivity to a chosen input and propagates through operations using the chain rule. section 4.3.2
JAX exposes this calculation through jax.jvp(). It returns the ordinary output and its tangent, or directional derivative. For several inputs, changing the tangent vector selects a direction; basis vectors recover Jacobian columns. section 4.3.2
A tangent through jvp()
For f(x1, x2) = x1*x2 + sin(x1*x2), evaluating x = (7.0, 2.0) with tangent (1.0, 0.0) produces the primal 14.990607 and tangent 2.2734745. The tangent is the derivative with respect to x1. section 4.3.2
Table 4.2. Forward mode autodiff example with f(x1, x2) = x1 × x2 + sin(x1 × x2) evaluated at the point (7.0, 2.0) with respect to the first variable x1 (view table figure)
| Forward primal trace |
Forward tangent (derivative) trace |
|---|---|
| v−1 = x1 = 7.0 (start) v0 = x2 = 2.0 |
v’−1 = x’1 = 1.0 (start) v’0 = x’2 = 0.0 |
| v1 = v−1 × v0 = 14.0 v2 = sin(v1) = 0.99 v3 = v1 + v2 = 14.99 |
v’1 = v’−1 × v0 + v−1 × v’0 = 1.0 × 2.0 + 7.0 × 0.0 v’2 = v’1 × cos(v1) = 2.0 × 0.137 v’3 = v’1 + v’2 = 2.0 + 0.274 |
| y = v3 = 14.99 (end) |
y’ = v’3 = 2.274 (end) |
Trace the primal and tangent values for each elementary operation in a forward-mode calculation. section 4.3.2
