A directional derivative measures a function’s slope at a point while moving in a chosen input direction. Unlike a partial derivative, which follows one coordinate axis, it accepts any direction vector. In JAX, pass that vector as the tangent to jax.jvp() and obtain the derivative in one forward pass without constructing the full Jacobian. section 4.3.2
Compute a directional derivative
For f(x1, x2) = x1x2 + sin(x1x2) at (7, 2), use jax.jvp(f, (7.0, 2.0), (1.0, 0.0)). The primal output is 14.990607 and the tangent output is 2.2734745, the slope in the positive x1 direction. Replacing the tangent with another vector measures a different direction. section 4.3.2
