Matrix-free nonlinear solves and implicit differentiation

Newton–Krylov root solve

newton_krylov solves \(F(x)=0\) without materializing the Jacobian. Each Newton update uses jax.linearize for the Jacobian-vector action and restarted flexible GMRES for the correction [KK04]. The state and residual may be arbitrary matching JAX PyTrees.

solution = sx.newton_krylov(
    residual,
    initial,
    precond=approximate_inverse,
    inner_product=weighted_product,
    rtol=1e-8,
    max_steps=10,
    linear_restart=20,
)

The nonlinear stopping test uses the true residual relative to the initial residual. Inspect both solution.converged and solution.linear_converged, as well as newton_iterations, linear_iterations, and residual_norm. inner_product= controls GMRES orthogonalization and supports weighted or distributed reductions; norm= may independently define the nonlinear residual norm.

For large PDE or kinetic systems, precond= should approximate the inverse Jacobian action using inexpensive problem structure. Jacobian-free products remove matrix storage, but do not by themselves make Krylov convergence mesh-independent [KK04].

Implicit differentiation computes derivatives of a converged equation rather than derivatives of the iterations used to solve it. This makes the gradient cost independent of the number of primal iterations and avoids storing their history [BBC+22, SB26].

Linear solve derivative

Let

\[ A(\theta)x(\theta)=b(\theta). \]

Differentiating gives

\[ A\,dx=db-(dA)x, \]

so

\[ dx=A^{-1}[db-(dA)x]. \]

For a reverse-mode cotangent \(\bar x\), solve the JAX linear-transpose system

\[ A^T\lambda=\bar x. \]

Then

\[ \bar b=\lambda, \qquad \bar A=-\lambda x^T. \]

The adjoint requires one transposed/adjoint solve regardless of how many steps the forward solver used.

linear_solve

def solver(operator, rhs):
    solution = sx.gmres(operator, rhs, rtol=1e-11)
    return solution.x, (
        solution.residual_norm,
        solution.iterations,
        solution.converged,
    )

x, diagnostics = sx.linear_solve(
    matvec,
    b,
    solver,
    transpose_matvec=adjoint_matvec,  # optional
    transpose_solver=solver,           # optional; defaults to solver
    has_aux=True,
)

Inputs:

• matvec: linear action closing over differentiable parameters;

• b: right-hand side;

• solver: callable (operator, rhs) -> x traceable by JAX, or (operator, rhs) -> (x, aux) when has_aux=True;

• transpose_matvec: optional explicit adjoint action;

• transpose_solver: optional dedicated solver for the transposed system;

• has_aux: return the forward solver’s auxiliary diagnostics with x.

If no explicit transpose is supplied, JAX constructs the linear transpose of matvec. By default the same solver callable is used for the transposed action. For nonsymmetric problems, pass transpose_solver when the adjoint needs a different initial guess, preconditioner, or stopping policy. With has_aux=True, both callbacks return (x, aux); the public result retains the forward diagnostics.

For complex programs, linear_solve follows JAX’s transpose/cotangent semantics rather than independently imposing a Hermitian-adjoint convention. Use a real scalar objective and validate directional derivatives, as in examples/18_complex_krylov_gradient.py.

Root derivative

Let \(x^*(\theta)\) satisfy

\[ f(x^*,\theta)=0. \]

The implicit function theorem gives

\[ \frac{dx^*}{d\theta} =-left(\frac{\partial f}{\partial x}\right)^{-1} \frac{\partial f}{\partial\theta}. \]

root = sx.root_solve(
    f,
    x0,
    solver=newton_solver,
    tangent_solve=tangent_linear_solver,
)

solver(f, x0) computes the primal root. tangent_solve(g, y) solves the linearized system \(g(v)=y\), where g is a Jacobian-vector product at the root. The default uses scalar division for scalar roots and materializes a dense Jacobian for vectors, so large roots should always supply a matrix-free tangent solver.

def tangent_solve(jvp, rhs):
    return sx.gmres(jvp, rhs, rtol=1e-10).x

PCG-specific wrapper

When the operator is HPD, pcg_linear_solve combines implicit gradients with forward PCGSolution diagnostics. It permits independent transpose tolerances, iteration limits, and preconditioners. See Preconditioned conjugate gradients.

Accuracy

The computed derivative is accurate only if:

1. the primal residual is sufficiently small;

2. the tangent/adjoint residual is sufficiently small;

3. the solution branch is locally differentiable;

4. the supplied transpose action matches the primal operator.

Gradient error is not controlled by primal tolerance alone. Tighten the adjoint solve and compare selected directional derivatives with central finite differences.

What implicit differentiation does not smooth

• branch selection in a multi-root problem;

• clipping, contact, or active-set changes;

• discontinuous remeshing or discrete model selection;

• a singular Jacobian at a bifurcation;

• an unconverged primal solve.

Implicit differentiation gives the local derivative of the selected regular branch. It should not be used to claim differentiability across a branch jump.

Comparison with unrolling

Strategy	Memory	Derivative meaning	Best use
differentiate iterations	grows with unrolled work unless rematerialized	finite algorithm	learned optimizers, deliberately truncated solves
implicit differentiation	one tangent/adjoint solve	converged equation	equilibrium, steady state, linear response
finite differences	one or more extra primal solves per direction	numerical directional derivative	validation, not high-dimensional gradients

API summary

• solvax.implicit.newton_krylov()

• solvax.implicit.linear_solve()

• solvax.implicit.root_solve()

• solvax.pcg.pcg_linear_solve()

Runnable counterparts: examples/06_implicit_differentiation.py, examples/17_pcg.py, and examples/18_complex_krylov_gradient.py.