Implicit solves and differentiation API

Implicit differentiation of linear solves and root finds.

For a parameterised linear system A(theta) x = b(theta), the implicit function theorem gives, without ever differentiating through the solver’s internal iterations,

dx = A^{-1} (db - dA x),

so the VJP of x = solve(A, b) against a cotangent xbar is obtained from a single transposed solve,

A^T lambda = xbar, bbar = lambda, Abar = -lambda x^T.

Likewise, for a root x* of f(x, theta) = 0,

dx*/dtheta = -(df/dx)^{-1} (df/dtheta),

which requires one linear solve against the Jacobian of f at the root. In both cases the forward solver is treated as a black box: it may iterate to arbitrary tolerance, use stopping criteria, restarts, preconditioning — none of that is differentiated. The adjoint costs exactly one additional (transposed / tangent) solve, independent of how many iterations the forward solve took.

These wrappers are thin layers over jax.lax.custom_linear_solve() and jax.lax.custom_root(), specialised to the “bring your own solver” pattern used throughout solvax (e.g. Krylov methods from solvax.krylov or the structured direct solves in solvax.direct).

References

  • D. A. Knoll & D. E. Keyes, Jacobian-free Newton–Krylov methods: a survey of approaches and applications, J. Comput. Phys. 193, 357 (2004).

  • M. Blondel et al., Efficient and Modular Implicit Differentiation, NeurIPS 2022, https://arxiv.org/abs/2105.15183.

  • C. S. Skene & K. J. Burns, Fast adjoints for kinetic solvers, https://arxiv.org/abs/2506.14792 — the one-extra-transposed-solve adjoint for iterative kinetic-equation solvers.

  • JAX documentation for jax.lax.custom_linear_solve and jax.lax.custom_root.

class solvax.implicit.NewtonKrylovSolution(x, residual_norm, newton_iterations, linear_iterations, converged, linear_converged)

Result of newton_krylov().

Parameters:
x: Any

Alias for field number 0

residual_norm: Array

Alias for field number 1

newton_iterations: Array

Alias for field number 2

linear_iterations: Array

Alias for field number 3

converged: Array

Alias for field number 4

linear_converged: Array

Alias for field number 5

solvax.implicit.newton_krylov(residual_fn, x0, *, precond=None, inner_product=None, norm=None, rtol=1e-08, atol=0.0, max_steps=20, linear_restart=30, linear_rtol=0.1, linear_atol=0.0, linear_max_restarts=10)

Solve a nonlinear system with matrix-free Newton–GMRES.

Jacobian-vector products are obtained from jax.linearize(); the Jacobian is never materialised. The nonlinear iteration stops when

norm(residual) <= max(atol, rtol * norm(initial_residual)).

Parameters:

  • residual_fn (Callable[[Any], Any]) – nonlinear residual with the same PyTree input/output structure.

  • x0 (Any) – initial iterate.

  • precond (Callable[[Any], Any] | None) – optional right preconditioner passed to GMRES.

  • inner_product (Callable[[Any, Any], Array] | None) – optional GMRES inner product; useful for weighted or distributed PyTrees.

  • norm (Callable[[Any], Array] | None) – optional nonlinear residual norm. Defaults to the norm induced by inner_product, or the Euclidean PyTree norm when both are omitted.

  • rtol (float) – nonlinear relative residual tolerance.

  • atol (float) – nonlinear absolute residual tolerance.

  • max_steps (int) – maximum Newton updates.

  • linear_restart (int) – GMRES Arnoldi cycle size.

  • linear_rtol (float) – GMRES relative residual tolerance for each Newton update.

  • linear_atol (float) – GMRES absolute residual tolerance.

  • linear_max_restarts (int) – maximum GMRES restart cycles per Newton update.

Returns:

A NewtonKrylovSolution with the final iterate and diagnostics.

Return type:

NewtonKrylovSolution

solvax.implicit.linear_solve(matvec, b, solver, *, transpose_matvec=None, transpose_solver=None, has_aux=False)

Differentiable linear solve with a user-supplied black-box solver.

Solves matvec(x) = b by calling solver(matvec, b), and registers an implicit-function-theorem VJP via jax.lax.custom_linear_solve(): gradients w.r.t. anything matvec or b close over cost exactly one additional transposed solve — the solver’s internal iterations are never differentiated.

Parameters:

  • matvec (Callable) – linear callable computing A @ x; parameters of A should be closed over so gradients can flow to them.

  • b (Array) – right-hand side.

  • solver (Callable) – callable solver(matvec, b) -> x, e.g. a lambda around a Krylov method or a dense solve. With has_aux=True, it must instead return (x, aux). It is treated as a black box (never differentiated through), but must be traceable by JAX.

  • transpose_matvec (Callable | None) – optional callable computing A^T @ y for the adjoint solve. If omitted, the transpose obtained from jax.linear_transpose of matvec is used.

  • transpose_solver (Callable | None) – optional solver for the transposed system. Defaults to solver. With has_aux=True, it must also return (x, aux).

  • has_aux (bool) – whether both solver callbacks return auxiliary data alongside the solution. The forward auxiliary data is returned with x.

Returns:

The solution x, differentiable w.r.t. the parameters of matvec and b. When has_aux=True, returns (x, aux) using the forward solver’s auxiliary data.

Return type:

Any

solvax.implicit.root_solve(f, x0, solver, *, tangent_solve=None)

Differentiable root find with a user-supplied black-box rootfinder.

Finds x with f(x) = 0 by calling solver(f, x0), and registers the implicit-function-theorem derivative via jax.lax.custom_root(): dx/dtheta = -(df/dx)^{-1} df/dtheta for any parameters theta closed over by f. The rootfinder’s internal iterations (e.g. a Newton loop with line search) are never differentiated.

The signature mirrors jax.lax.custom_root(); parameters of f beyond x should be closed over so gradients can flow to them.

Parameters:

  • f (Callable) – function whose root is sought, f(x) -> residual with the same shape as x.

  • x0 (Array) – initial guess passed to solver.

  • solver (Callable) – callable solver(f, x0) -> x_root, e.g. a Newton loop. Treated as a black box (never differentiated through), but must be traceable by JAX.

  • tangent_solve (Callable | None) – optional callable tangent_solve(g, y) -> x solving the linearised system g(x) = y, where g is the Jacobian of f at the root as a linear map. Defaults to a scalar division for scalar problems and a dense jnp.linalg.solve() of the materialised Jacobian for small vector systems.

Returns:

The root x, differentiable w.r.t. the parameters closed over by f.

Return type:

Array