liveBook · Manning

Vector-Jacobian product

As explained in Deep Learning with JAX

The product of an incoming vector with the Jacobian of a function, computed to propagate sensitivities efficiently.

A vector-Jacobian product (VJP) propagates a chosen output sensitivity backward through a function. JAX's reverse-mode transform returns both the function value and a callable that maps an output cotangent to input cotangents. section 4.3.3

Computing a VJP with JAX

For f(x1, x2) = x1x2 + sin(x1x2), use p, vjp_func = jax.vjp(f, 7.0, 2.0). The primal output is 14.990607; vjp_func(1.0) returns approximately (2.2734745, 7.9571605), matching reverse-mode gradients. section 4.3.3

VJPs are useful for wide Jacobians: applying different output cotangents can recover Jacobian rows without materializing the full Jacobian. Reverse mode saves passes when inputs greatly outnumber outputs, but requires more storage. section 4.3.3

How Deep Learning with JAX covers it

How deep it goes
Explains the returned callable, cotangent shapes, Jacobian-row recovery, and storage tradeoff without deriving transformation internals.
Examples
Demonstrates scalar reverse differentiation and recovers two Jacobian rows for a two-output function.
What it assumes
Assumes functions, Jacobians, derivatives, and cotangent vectors; demonstrates the JAX call pattern from concrete examples.
What you can do after
Apply it via JAX's vjp() to obtain primal outputs, backpropagate chosen cotangents, and inspect Jacobian rows.
How it connects
Connects reverse-mode autodiff, backpropagation, Jacobians, Jacobian-vector products, cotangents, and gradients.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.