quadax.LeibnizAdjoint

class quadax.LeibnizAdjoint(options: dict = <factory>, options_fwd: dict = <factory>, options_rev: dict = <factory>)Source

Differentiate by the Leibniz rule, either mode, with its own error control.

The derivative is evaluated with its own adaptive solve, so it gets its own error control rather than inheriting the subdivision chosen for the integral; see Adjoints for when that is worth paying for.

Because each mode picks its own subdivision, forward and reverse results agree to quadrature accuracy rather than exactly.

Parameters:

• options (dict, optional) – Options for the derivative’s own solve, replacing the ones the quadrature routine was called with. Default is to solve for the derivative with the same options as the integral.

• options_fwd (dict, optional) – The same, for one direction alone, taking precedence over options. Which vector is being measured differs between the two directions, so norm is the option most likely to want this; see the note below.

• options_rev (dict, optional) – The same, for one direction alone, taking precedence over options. Which vector is being measured differs between the two directions, so norm is the option most likely to want this; see the note below.

Notes

The two directions of the derivative do not integrate the same vector, which is what options_fwd and options_rev are for.

Forward mode integrates the tangent of the integrand, a vector of the integrand’s own shape. Reverse mode integrates the cotangent of the arguments being differentiated, raveled into a single flat vector. Its length is the total number of differentiated parameter components, which has nothing to do with the integrand’s shape, and its entries are parameters rather than components of the integral. They appear in the order the quadrature’s own arguments do: the limits, args, and whatever the integrand closes over. Each contributing size entries if it is being differentiated and none at all if it is not. Differentiating a two point interval and a length three args[0] gives a vector of five with the limits first; differentiating args[0] alone gives a vector of three.

A norm weighted per parameter is written against that layout, and belongs in options_rev so that it is never handed a tangent instead, which the two vectors being the same length by coincidence would otherwise hide.

When differentiating a moving jump or singularity, mark the jump or singularity with a breakpoint, and build that breakpoint from the same parameter that positions the feature. Marking a feature that is genuinely there is never worse than leaving it unmarked, and for derivatives it is frequently the difference between a correct answer and a silently wrong one. The one way marking can hurt is marking something that is not there, which is the third case below.

Differentiating a jump gives a delta, which no quadrature of the integrand’s tangent can represent, so it is recovered from the motion of the breakpoint instead, and that works only when the breakpoint moves with the discontinuity:

step = lambda t, z: jnp.where(t > z[0], 1.0, 0.0)   # jumps at t = z[0]

# correct: one parameter `s` positions both the jump and the breakpoint
f = lambda s: quadgk(step, jnp.array([-1.0, s, 1.0]), (jnp.array([s]),))[0]
jax.grad(f)(0.3)            # -1.0

Two versions that look equivalent are not. Both return the same primal value, 0.7, and both give zero where the derivative is -1:

# WRONG: the jump is unmarked, so nothing tracks it
f = lambda s: quadgk(step, jnp.array([-1.0, 1.0]), (jnp.array([s]),))[0]

# WRONG: marked, but with a constant, which carries no derivative. Tempting,
# because it does help the primal value: it cuts the subdivision from 13
# sub-intervals to 2.
f = lambda s: quadgk(step, jnp.array([-1.0, 0.3, 1.0]), (jnp.array([s]),))[0]

Splitting the feature across two parameters is what the requirement rules out. The derivative with respect to a breakpoint on its own is not a well posed question: the integral is the same whatever the mesh is cut at, so the breakpoint only means anything in combination with the integrand it is marking. Ask for it anyway, by differentiating with respect to [breakpoint, jump location] at [0.3, 0.3], and the answer is [-1, 0]. The total over the two is right, and how it is divided between them is an artifact of having written one feature as two parameters, not a property of the quadrature.

A moving singularity needs the same treatment, for the same reason: unmarked, it slides across a subdivision that does not move with it.

Only a feature that is steep but continuous needs none of this: a sharp tanh differentiates correctly unmarked under either adjoint. Marking one anyway is harmless, the breakpoint then only helps the subdivision, and contributes nothing to the derivative since the integrand has the same limit from both sides of it.

Methods

quadrature(ops, interval, args, consts, ...)

Evaluate the quadrature, differentiating it by the Leibniz rule.

Attributes

options
options_fwd
options_rev