A Jacobian-vector product (JVP) applies a function’s Jacobian to an input tangent vector without constructing the full Jacobian. In JAX, jax.jvp() performs this calculation in one forward-mode pass. section 4.3.2 section 4.4
Pass the function, its primal values, and matching tangent values. The result is a (primals_out, tangents_out) pair: the ordinary function output and its directional derivative. Choose a direction such as [1, 1, 1] to probe the combined change along all three inputs. section 4.3.2
A small JAX calculation
For f2(x) = [x0² + x1² − x1x2, x0² − x1² + 3x0x2], evaluate x = [3, 4, 5] with tangent v = [1, 1, 1]: p, t = jax.jvp(f2, (x,), (v,)) produces p = [5, 38] and t = [5, 22]. Repeating the call with unit tangents recovers individual Jacobian columns, but requires separate forward passes. 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 x1 × x2 + sin(x1 × x2), evaluated at (7, 2), to see forward-mode propagation produce a derivative of about 2.274. section 4.3.2
