quadax.DirectAdjoint

class quadax.DirectAdjoint(checkpoint: bool = False, chunk_size: int = 8)Source

Differentiate the quadrature exactly, reusing the converged subdivision.

Commonly called “discretize then optimize”: the quadrature is discretized first, and the derivative is then taken of that discretization.

This is the default, and is the cheaper option for a cheap integrand in either mode. When one evaluation of the integrand is expensive, LeibnizAdjoint can be an order of magnitude faster or more; see Adjoints for the trade-offs.

It works by running the primal solve recording the final adaptive mesh, and then using the same mesh (with corrections when differentiating the interval itself) to integrate the derivative of the integrand.

The derivative therefore inherits the subdivision chosen for the integral, and no error control of its own is paid for. That is the reason to reach for a Leibniz adjoint when the derivative needs resolving that the integral did not pay for, and the reason a derivative costs roughly what the converged subdivision costs however generous max_ninter was.

Methods with no subdivision to reuse (Romberg) are handled the same way in spirit: the number of Richardson levels the solve settled on is frozen instead of a mesh. There the fixed discretization still contains a fori_loop with dynamic bounds that JAX cannot reverse differentiate, so the derivative is routed through a custom primitive that supplies the two directions explicitly. Both modes work either way.

Parameters:

• checkpoint (bool) – Whether to recompute the quadrature on each block of sub-intervals during the backward pass rather than storing it, off by default. Without it reverse mode keeps the integrand’s value at every node of every sub-interval, which dominates its memory and grows with max_ninter however few sub-intervals are really used. Turning checkpoint on can reduce memory by 3x to 30x depending on how large max_ninter is at the expense of additional integrand evaluations. No effect in forward mode.

• chunk_size (int) – How many sub-intervals of the frozen subdivision to evaluate at once: vmap within a chunk, scan across chunks. Trades peak memory against speed, and is the reverse-mode counterpart of the quadrature routines’ batch_size, which bounds the work within one sub-interval. The two multiply, so a gradient evaluates the integrand at up to chunk_size times batch_size points at a time. Lower it when a derivative runs out of memory and checkpoint was not enough; raise it when the subdivision is small and the scan is pure overhead.

Notes

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 and differentiate it on the converged subdivision.

Attributes

checkpoint
chunk_size