| name | jet-differentiation |
| description | Uses second-order Jets (value, gradient, Hessian triples) for automatic differentiation in the local_coordinates JAX library, including pushforward through smooth maps, change of coordinates, and the @jet_decorator pattern. Use when the user works with Jets, second-order Taylor expansions, JVPs and HVPs, or needs to propagate derivatives through a composition of maps. |
Jet differentiation
Use this skill when working with second-order Jets. A Jet J[F]_p = (F_p, dF_p/dx, d^2F_p/dx^2) captures the second-order Taylor expansion of F at p.
Typical workflow
function_to_jet(f, x) materializes a Jet at x using jax.jacrev for the gradient and jax.jacfwd(jax.jacrev) for the Hessian.- Decorate a function with
@jet_decorator to lift it so it operates on Jet inputs, propagating derivatives automatically via JVPs.
- All Jets passed through a decorated function must share the same coordinate dimension. PyTree-valued Jets are supported as long as value, gradient, and Hessian share structure.
Transformation laws
Under a change of coordinates x -> z, the gradient transforms with the inverse Jacobian and the Hessian picks up a correction term involving the Hessian of the coordinate change:
d^2 F / dz^i dz^j = (dx^b/dz^j) [d^2 F / dx^d dx^b - (d^2 z^c / dx^b dx^d)(dx^a / dz^c)(dF / dx^a)] (dx^d / dz^i)
Under composition G = T o F, the Hessian of G has a curvature term from the nonlinearity of T and a transport term from the Hessian of F.
Invariants
- Value preservation under coordinate change.
- Chain rule consistency for both gradient and Hessian.
- Hessian symmetry across the two derivative indices.
- Coordinate round-trip recovers the original Jet.
See references/REFERENCE.md for the full derivation and scripts/sanity_check.py for the verification suite.
