04. Newton’s method: a loop that stops when it has converged

Data-dependent control flow: differentiate through a while loop with a convergence test, and through the in-place Thomas sweeps inside it.

The problem

Solve the nonlinear boundary-value problem \[F(u, \theta) = A u + u^3 - \theta = 0, \qquad A = \tfrac{1}{h^2}\operatorname{tridiag}(-1, 2, -1),\] by Newton’s method until the step is small, \(\lVert\delta\rVert_\infty \le 10^{-10}\lVert u\rVert_\infty\), and differentiate \(\mathcal{L}(\theta) = \lVert u^\ast(\theta) - u_{\text{obs}}\rVert^2\) with respect to \(\theta\). Each Newton step solves a tridiagonal system with the Thomas algorithm, a sequential in-place sweep.

Two irregular patterns:

  • Data-dependent control flow. The number of Newton steps depends on \(\theta\) (3–4 here).
  • Sequential in-place update. The Thomas forward elimination and back substitution.

Code

"Residual F(u, θ) = A u + u³ − θ with zero Dirichlet ends, written into F."
function residual!(F, u, θ)
    n = length(u)
    ih2 = (n + 1)^2
    for i in 1:n
        l = i > 1 ? u[i-1] : zero(eltype(u))
        r = i < n ? u[i+1] : zero(eltype(u))
        F[i] = ih2 * (2u[i] - l - r) + u[i]^3 - θ[i]
    end
end

"Solve J δ = F in place (δ overwrites F), J = A + 3 diag(u²): the Thomas algorithm."
function thomas!(F, c, u)
    n = length(u)
    ih2 = (n + 1)^2
    off = -ih2
    b = 2ih2 + 3u[1]^2
    c[1] = off / b
    F[1] = F[1] / b
    for i in 2:n                                  # forward elimination
        b = 2ih2 + 3u[i]^2 - off * c[i-1]
        c[i] = off / b
        F[i] = (F[i] - off * F[i-1]) / b
    end
    for i in n-1:-1:1                             # back substitution
        F[i] -= c[i] * F[i+1]
    end
end

"Newton until the step is small; returns the solution and the iteration count."
function newton_solve(θ::AbstractVector{T}) where {T}
    n = length(θ)
    u, F, c = zeros(T, n), zeros(T, n), zeros(T, n)
    it, converged = 0, false
    while !converged && it < MAXITER              # data-dependent trip count
        residual!(F, u, θ)
        thomas!(F, c, u)                          # F now holds the Newton step δ
        u .-= F
        it += 1
        converged = maximum(abs, F) <= TOL * maximum(abs, u)
    end
    return u, it
end

newton_loss(θ, uobs) = sum(abs2, first(newton_solve(θ)) .- uobs)
newton_loss (generic function with 1 method)
Source: 04. Newton solver — Enzyme.jl
"Reverse mode straight through the while loop, the Newton steps and the Thomas sweeps."
function newton_gradient(θ, uobs)
    dθ = zero(θ)
    Enzyme.autodiff(Reverse, newton_loss, Active, Duplicated(θ, dθ), Const(uobs))
    return dθ
end

g = newton_gradient(θ0, uobs0)
@printf "Newton iterations: %d   loss = %.12f   ‖∇loss‖ = %.6f\n" last(newton_solve(θ0)) newton_loss(θ0, uobs0) sqrt(sum(abs2, g))
Newton iterations: 4   loss = 17.067028390381   ‖∇loss‖ = 0.737983

Implicit-function alternative:

# The alternative a numerical analyst would choose: at the solution, F(u*, θ) = 0, so
# dL/dθ = λ with Jᵀλ = ∂L/∂u (J is symmetric here). One tridiagonal solve, no tape.
function newton_gradient_implicit(θ, uobs)
    u, _ = newton_solve(θ)
    λ, c = 2 .* (u .- uobs), similar(u)
    thomas!(λ, c, u)
    return λ
end

g_imp = newton_gradient_implicit(θ0, uobs0)
@printf "unrolled vs implicit: max rel. difference %.1e\n" maximum(abs, g - g_imp) / maximum(abs, g)
unrolled vs implicit: max rel. difference 2.6e-14
def residual(u, theta):
    n = u.shape[0]
    ih2 = (n + 1) ** 2
    up = jnp.pad(u, 1)
    return ih2 * (2 * u - up[:-2] - up[2:]) + u**3 - theta

def newton_step(u, theta):
    n = u.shape[0]
    ih2 = (n + 1) ** 2
    lower = jnp.full(n, -ih2, u.dtype).at[0].set(0.0)
    upper = jnp.full(n, -ih2, u.dtype).at[-1].set(0.0)
    diag = 2 * ih2 + 3 * u**2
    delta = jax.lax.linalg.tridiagonal_solve(lower, diag, upper, residual(u, theta)[:, None])[:, 0]
    return u - delta

# The direct translation: Newton until the step is small.
def loss_while(theta, uobs):
    def cond(state):
        u, step, it = state
        return (step > TOL * jnp.max(jnp.abs(u))) & (it < MAXITER)
    def body(state):
        u, _, it = state
        u_new = newton_step(u, theta)
        return u_new, jnp.max(jnp.abs(u_new - u)), it + 1
    u, _, _ = jax.lax.while_loop(cond, body, (jnp.zeros_like(theta), jnp.inf, 0))
    return jnp.sum((u - uobs) ** 2)

# Workaround 1: a fixed number of Newton steps, enough for the worst case.
def loss_fixed(theta, uobs, iters=MAXITER):
    u = jax.lax.fori_loop(0, iters, lambda _, u: newton_step(u, theta), jnp.zeros_like(theta))
    return jnp.sum((u - uobs) ** 2)

# Workaround 2 (idiomatic): Optimistix root finding, differentiated implicitly.
def loss_implicit(theta, uobs):
    # Optimistix stops when the step is small *and* |F| < atol. Rounding in F grows like
    # (n+1)²·eps, so we root-find on h²·F instead: same root, same Newton steps, same
    # implicit gradient, but a residual whose rounding error is ~eps·|u|.
    def scaled_residual(u, theta):
        return residual(u, theta) / (u.shape[0] + 1) ** 2
    solver = optx.Newton(rtol=TOL, atol=1e-12, linear_solver=lx.Tridiagonal())
    sol = optx.root_find(scaled_residual, solver, jnp.zeros_like(theta), theta,
                         tags=frozenset({lx.tridiagonal_tag}), max_steps=MAXITER)
    return jnp.sum((sol.value - uobs) ** 2)
Source: 04. Newton solver — JAX
try:
    jax.grad(loss_while)(theta0, uobs0)
except ValueError as err:
    print("jax.grad(loss_while):", str(err).splitlines()[0])

grad_fixed = jax.jit(jax.grad(loss_fixed))
grad_implicit = jax.jit(jax.grad(loss_implicit))
g = grad_implicit(theta0, uobs0)
g_fixed = grad_fixed(theta0, uobs0)
print(f"loss = {float(loss_implicit(theta0, uobs0)):.12f}   ‖∇loss‖ = {float(jnp.linalg.norm(g)):.6f}")
print(f"fixed-iteration vs implicit: max rel. difference {float(jnp.max(jnp.abs(g - g_fixed)) / jnp.max(jnp.abs(g))):.1e}")
jax.grad(loss_while): Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.
loss = 17.067028390381   ‖∇loss‖ = 0.737983
fixed-iteration vs implicit: max rel. difference 4.2e-14

The direct JAX translation, jax.grad of a lax.while_loop, is rejected: reverse mode needs a static trip count. JAX offers two ways out: run a fixed, worst-case number of steps, or switch to implicit differentiation (Optimistix). Enzyme differentiates the loop as written, and the implicit adjoint is still available as a two-line alternative.

Same gradient

Enzyme differentiates the Newton iterates and Optimistix applies the implicit-function theorem at its own stopping point. Newton converges quadratically, so both land on the solution to rounding and the gradients agree closely.

implementation ‖g‖₂ max rel. diff vs enzyme
enzyme 0.737983 0.0e+00
jax 0.737983 1.4e-14

Timings

Float64, median of 3–10 runs after warm-up. “Forward” is the converged Newton solve; for JAX that is the while_loop version.

What this shows

At \(n = 10^6\) the gradient takes:

gradient
Enzyme through the loop (CPU) 130.3 ms
Enzyme, implicit adjoint (CPU) 69.1 ms
JAX, fixed 50 steps (CPU) 3.32 s
JAX, Optimistix implicit (CPU) 185.4 ms
JAX, fixed 50 steps (GPU) 63.4 ms
JAX, Optimistix implicit (GPU) 3.5 ms

Differentiating through the iterations stores every Newton step for the reverse pass. That costs time and memory (at \(n = 10^6\), Enzyme allocates 69 MiB unrolled against 38 MiB implicit), and it is the reason JAX refuses dynamic loops. Implicit differentiation is usually the better method anyway, and the table shows that.

The fixed-count workaround is expensive: it has to run the worst case (here MAXITER = 50 steps) although Newton converges in 3–4, and reverse mode then stores all 50. Choosing a smaller count means guessing how many steps every future input will need. The difference is who decides. With Enzyme it is a choice: the loop as written works, and the implicit adjoint is a few lines when you want it. With JAX you are forced into implicit differentiation or fixed iteration counts, and for a solver without a clean fixed-point structure (adaptive time stepping, line searches, event detection) there may be no implicit formulation to fall back on.

Where JAX wins / what we did not do

Note

On the GPU, JAX with Optimistix’s implicit differentiation is by far the fastest: 3.5 ms at \(n = 10^6\), using XLA’s GPU tridiagonal solver. Enzyme’s version here is single-threaded CPU code.

Using Optimistix did take care: its default stopping test requires \(|F| <\) atol, which rounding makes unreachable for this problem at large \(n\), and its alternative step-based test reports “diverged” once the steps stall at rounding level. The fix, root-finding on the rescaled residual \(h^2 F\), is in the JAX code above.

  • No Enzyme GPU version. The Thomas algorithm is sequential, so a GPU version would need a different tridiagonal solver (cyclic reduction) on both sides.
  • No adaptive time-stepping example (Diffrax), which was in the original plan as a stretch goal.
  • No Reactant column on this page yet.