Autodiff products

The question this method answers¶

How does a scientific program’s output change under a small, specified change to its input, and how can that local change be computed without forming a dense Jacobian? Autodiff products answer this question for the program that JAX actually executes. Begin with What is a derivative? if derivatives as local linear maps are new.

Before computation: what should be true?¶

The function must accept and return floating JAX arrays with shapes compatible with the proposed tangent or cotangent. More importantly, the executed branch must represent the scientific perturbation being claimed. A derivative through a clip, discrete index, or branch transition can be finite yet answer the wrong question.

Define the mathematical objects¶

Let $f:\mathbb{R}^n\rightarrow\mathbb{R}^m$ be differentiable at $x$ . Its derivative $D f(x)$ is the linear map that gives the first-order response

f(x+\epsilon v)=f(x)+\epsilon D f(x)[v]+o(\epsilon),

(1)

where $v\in\mathbb{R}^n$ is a tangent direction and $o(\epsilon)/\epsilon\to0$ . The Jacobian $J\in\mathbb{R}^{m\times n}$ is a coordinate representation of that map. A cotangent $w\in\mathbb{R}^m$ weights output directions. For a scalar $f:\mathbb{R}^n\to\mathbb{R}$ , the Hessian $H=\nabla^2 f(x)\in\mathbb{R}^{n\times n}$ describes local curvature.

The data representation matters: parameter arrays and scientific PyTrees are discussed in PyTrees as scientific state.

Derive the method¶

The JVP pushes the input direction $v$ through the derivative:

\operatorname{JVP}(f,x,v)=D f(x)[v]=Jv.

(2)

The VJP pulls the output cotangent $w$ back to input space:

\operatorname{VJP}(f,x,w)=D f(x)^{\mathsf T}[w]=J^{\mathsf T}w.

(3)

These are adjoint operations. Their defining scalar identity is

w^\mathsf{T} D f(x)[v] =w^\mathsf{T}Jv =v^\mathsf{T}J^\mathsf{T}w.

(4)

For scalar $f$ , applying a JVP to the gradient avoids materializing $H$ :

H v = D(\nabla f)(x)[v].

(5)

For residuals $r(x)\in\mathbb{R}^m$ , the least-squares objective $F(x)=\tfrac12 r(x)^\mathsf{T}r(x)$ has the Gauss-Newton curvature approximation $J_r^\mathsf{T}J_r$ . Its product is computed as one JVP followed by one VJP. For per-example score vectors $s_i$ , the empirical Fisher-style product is $N^{-1}\sum_i s_i(s_i^\mathsf{T}v)$ .

What the algorithm actually does¶

jvp delegates to jax.jvp and returns (f(x), Jv). vjp constructs JAX’s pullback and returns (f(x), J.T @ w). The product-only aliases discard the primal value. hvp applies jax.jvp to jax.grad(f). gauss_newton_product chains the module’s JVP and VJP helpers. empirical_fisher_product vmaps a two-argument score function over the leading data axis, stacks the scores, and applies the mean outer-product matrix without constructing that matrix.

No helper sanitizes non-finite values, checks scientific units, or changes JAX’s dtype rules. Shape, tracing, and dtype errors propagate.

What JAX differentiates¶

JAX differentiates the finite program represented by f along the supplied direction. JVPs use forward-mode linearization; VJPs use a reverse-mode pullback. Curvature products differentiate the executed gradient or residual program, including its smooth branches and any local saturation.

Using it in Jaxstro¶

Use the owner-qualified module so the runtime boundary is explicit:

import jax.numpy as jnp

from jaxstro.numerics.autodiff import jvp, vjp


def model(x):
    return jnp.array([x[0] ** 2 + x[1], jnp.sin(x[1])])


x = jnp.array([2.0, 0.5])
v = jnp.array([0.1, -0.2])
w = jnp.array([1.0, 3.0])
value, pushed = jvp(model, x, v)
_, pulled = vjp(model, x, w)

Here x and v have shape (2,), value and w have shape (2,), pushed has the output shape, and pulled has the input shape. hvp requires a scalar output. The current empirical Fisher helper assumes vector parameters and per-example vector scores compatible with ordinary matrix products.

How to audit the result¶

Choose a point away from known nonsmooth boundaries. Compare $Jv$ with the central directional finite difference

\frac{f(x+hv)-f(x-hv)}{2h},

(6)

then repeat over a decreasing sequence of $h$ values to separate truncation error from roundoff. Check the adjoint identity in (4) with independent $v$ and $w$ . For an HVP, finite-difference the gradient, not the original scalar function. Record dtypes, units, step sizes, absolute and relative disagreements, and whether the executed branch stayed fixed.

The package-wide audit vocabulary and executable evidence are in Validation methods.

Where the claim stops¶

These helpers reduce the cost and clarify the spelling of derivative products. They do not prove differentiability, condition a model, choose meaningful directions, certify a Hessian, or supply inference semantics. Dense Jacobian parity on a toy problem is implementation evidence, not scientific validation of a downstream model.