Solvers & preconditioners

jno.solve and jno.precond are slots: you hand one to fem.solve(...) or fdm.solve(...) and it replaces the default without changing a single term of the weak form. The same slot objects work for both front doors.

Between "accept the default" and "write a full solve_fn" sits the slot API: the solver factorises into four orthogonal slots, each a configured callable (never a string) from the jno.solve / jno.precond namespaces — or your own with the same contract. Every None keeps today's default; solve_fn= stays the total override (passing both is an error).

u = fem.solve(
    x0        = u_guess,                 # warm start — see "What a warm start buys" below
    nonlinear = jno.solve.newton(),      # linearization driver (nonlinear problems)
    linear    = jno.solve.gmres(),       # inner linear solve: lu / dense / cg / bicgstab / gmres
    precond   = jno.precond.jacobi(),    # v -> M⁻¹v spec, materialized against the assembled A
)

Choosing a linear solver

Everything shipped is pure JAX — jit/vmap-native, differentiable (Krylov on lax.custom_linear_solve, Newton on lax.custom_root), receives the assembler's BCOO operator directly (no densification), and composes with the periodic reduction and every parametric/inverse path. Pick by structure:

structure	solver	notes
SPD (Poisson, elasticity, mass)	cg	cheapest per iteration
non-symmetric (advection, SUPG)	bicgstab / gmres	bicgstab with jacobi() is the default
iterative preconditioner (inner Krylov, block/Schur)	fgmres	flexible right preconditioning — Saad, SISSC 14(2), 1993
symmetric indefinite (Stokes/Biot saddle, biharmonic)	minres	monotone residual, O(1) memory — Paige & Saunders, SINUM 12(4), 1975
complex-symmetric (A = Aᵀ, not Aᴴ: time-harmonic eddy-current/A–V, Helmholtz, RCWA)	cocg	CG's short recurrence over the bilinear form xᵀy — O(n) memory and one matvec per step where gmres needs a (restart, n) basis. Only valid when the operator really is complex-symmetric; it can break down where CG cannot — van der Vorst & Melissen, IEEE Trans. Magn. 26(2), 1990
SPD, batched/GPU-heavy	chebyshev	inner-product free (no reductions) — Golub & Varga 1961
indefinite	lu	sparse-direct; vmaps on every backend — host/cudss/pardiso factor once for a batch against one matrix, the default device backend factors each system (jno.setup(lu_stack=k) stacks k per call)
cuSolver refuses it, or is slow	lu(backend="host")	SuperLU on the host, driven from the device — same answer, same gradients. Often faster than cuSolver, and runs meshes it rejects
shift-invert eigs, or a constant-operator transient	lu(backend="cudss")	NVIDIA cuDSS. Fastest per solve — caches the symbolic plan on the sparsity, so it survives a change of values. Needs the optional stack
a Newton loop (or no GPU / a factorization too big for device memory)	lu(backend="pardiso")	Intel MKL PARDISO, multithreaded CPU. Fastest factorization — a Newton step reuses the analysis. x86-64
a Newton march (the tangent changes every iteration)	lu(backend="host", reuse=False)	the host factorization cache keys on the operator, so a never-repeating tangent fills it and leaks a factorization a step; reuse=False drops each one
small systems / coarse blocks	dense	LAPACK, vmap-native

jNO warns you when this applies

jno.fem detects a saddle system structurally at build time — a field whose own test function never meets its own trial function contributes no diagonal block — and fem.solve() warns, naming the field, when it is about to use the matrix-free default on one. It warns rather than refuses: a 2-D saddle of moderate size does solve acceptably that way. Passing linear= or precond= silences it, since that is the deliberate choice it asks for. The detection is structural, so it holds in every mode with no tangent to assemble, and it fires under jit/vmap/grad — which matters, because the runtime residual guard needs a concrete residual and steps aside on a tracer, leaving a transformed solve otherwise unguarded.

Warm starts — what x0= actually buys

x0= is free to pass and easy to misjudge, so here is the rule, measured. A Krylov solve has to cover the distance from its initial residual down to rtol. A warm start removes the decades it already starts below, so the fraction of iterations saved is a ratio of logarithms — not a ratio of errors:

\[\text{saved} \;\approx\; \frac{\log_{10}(1/\varepsilon)}{\log_{10}(1/\texttt{rtol})}\]

with ε the relative error of the guess. At rtol = 1e-10, a guess good to 1 % removes one fifth of the work — and a guess good to 10 % removes one tenth, which is usually not worth arranging.

Jacobi-CG, 2-D diffusion with a 10× conductivity blob, 791 DOFs, rtol = 1e-10

Cold (x0 = 0) is 85 iterations. Two initial errors of the same norm were compared: a smooth one (a low-frequency field) and a rough one (white noise).

rel. error ε	smooth x0	rough x0	law predicts
3e-01	78	93	81
1e-01	74	90	76
1e-02	67	85	68
1e-03	61	76	59
1e-04	52	67	51
1e-06	35	53	34

The law tracks the smooth column to within 1–3 iterations across five decades. The rough column does not follow it at all: a rough guess accurate to 1 % costs exactly what a cold start costs, and a rough guess at 30 % is worse than starting from zero.

The spectrum of the initial error matters, not its norm. Krylov convergence is governed by how the error is distributed over the spectrum (Greenbaum, Iterative Methods for Solving Linear Systems, SIAM 1997, ch. 3), and a guess that is pointwise close but rough injects high-frequency error the method must then remove. This is why the previous solve in a sweep or a time step is such a good warm start — it is smooth, and it is wrong in the same modes the operator is about to correct.

A neural surrogate is a poor warm start, measured

The tempting idea — train a network to predict the solution, feed it as x0 — does not clear the bar. On the sweep above, against a cold start: an MLP predicting DOF values directly came in at −2 % (i.e. slower than cold, its output being spectrally rough), and predicting POD coefficients instead — whose output is smooth by construction, being a combination of smooth modes — recovered only +5 %. Both lost to the trivial baseline of reusing the nearest stored snapshot (+9 %), which costs nothing to build.

The arithmetic is why, and it is not an engineering problem: at a realistic surrogate accuracy of 1e-2 to 1e-3 the law caps the prize at 20–30 %, and collecting it needs an output that is smooth as well as accurate. Reach for x0= where a nearby solve already exists — a continuation sweep (jno.solve.continuation, which warm-starts for you), a transient step (which warm-starts automatically, and rejects x0= for that reason), or a Newton loop. Those are smooth by construction and free.

A learned subspace is worth far more than a learned point

The failure above is a failure of regression, not of learning. Fit a network to predict a point and it optimises the wrong norm; CG minimises ‖e‖_A, and an L2 fit is blind to how the error sits on the spectrum. Give it a subspace instead and there are two much better things to do with it, both already expressible with the slots on this page.

1. Solve in the subspace rather than regressing into it. The A-orthogonal projection U(UᵀAU)⁻¹Uᵀb is the provably best starting point from span(U) — and it is exactly what fem.solve(basis=U) computes, certificate included.

2. Put the subspace in the preconditioner, as a coarse-space correction M⁻¹ = diag⁻¹ + U(UᵀAU)⁻¹Uᵀ (Nicolaides, SINUM 24(2), 1987, 355; Frank & Vuik, SISC 23(2), 2001, 442). This needs no new API — precond= takes a bare ctx -> applier:

def coarse_jacobi(U):                       # Jacobi + a coarse solve on span(U)
    def spec(ctx):
        AU   = jnp.stack([ctx.A @ U[:, j] for j in range(U.shape[1])], axis=1)
        Einv = jnp.linalg.inv(U.T @ AU)
        dinv = 1.0 / ctx.diag()
        return lambda v: dinv * v + U @ (Einv @ (U.T @ v))
    return spec

u = fem.solve(linear=jno.solve.cg(), precond=coarse_jacobi(U), basis=None)

Same problem, same rank-8 POD subspace, iterations against a cold Jacobi-CG

what the subspace is used for	saved
MLP regressing DOF values → x0	−2.7 %
MLP regressing POD coefficients → x0	+3.9 %
nearest stored snapshot → x0	+8.6 %
basis=U (A-optimal projection) → x0	+18.0 %
coarse-space preconditioner	+27.5 %
both together	+39.8 %

The L2-fitted network's energy-norm error was 17–77× larger than the A-optimal projection from the very same subspace — the whole gap between the first rows and the fourth is the wrong norm, not a weak network.

And the two are different in kind — only one improves with size

h	DOFs	cold	x0 saves	coarse-space saves
0.060	377	60	21.7 %	21.7 %
0.040	791	86	17.4 %	25.6 %
0.030	1394	113	17.7 %	28.3 %
0.022	2522	151	17.2 %	29.8 %

A warm start is a constant offset — it removes decades of residual and the log law fixes the fraction, so it stays flat under refinement. A coarse space removes the slow modes from the spectrum, so its benefit grows with the problem. Prefer it wherever the subspace is worth building at all.

Scope, since this was measured on one problem class: SPD + CG, a rank-8 POD subspace from snapshots, a 10× conductivity contrast. Non-symmetric and saddle-point systems are not covered here, the rank/cost trade-off is unexplored (each application costs two n × k products and a k × k solve), and POD alone already delivers the table — a learned subspace has to beat POD to earn its place, which is not shown.

Two that are not Ax = b on a square, definite operator

minres and lstsq cover the cases the table's other rows cannot, and both take a jno.solve.LinearOperator (or a BCOO, or a dense array) directly — no FEM required:

op = jno.solve.LinearOperator(K)          # K symmetric but INDEFINITE — cg cannot be used
x  = jno.solve.minres(tol=1e-12)(op, b)   # rel-resid 7.3e-13

xl = jno.solve.lstsq(A, b, atol=1e-12, btol=1e-12)   # A RECTANGULAR (80 x 30), min ‖Ax − b‖
xd = jno.solve.lstsq(A, b, damp=0.5)                 # + Tikhonov damp²‖x‖² — ill-posed inverses

Checked against dense oracles

minres on a spectrum straddling zero (eigenvalues in [-3, -0.5] ∪ [0.5, 3]) reaches a 7.3e-13 relative residual, where cg has no convergence theory at all. lstsq matches numpy.linalg.lstsq to 7.5e-13, and with damp=0.5 matches the Tikhonov normal equations (AᵀA + damp²I)x = Aᵀb to 5.1e-12.

Direct backends

Four lu backends, and the choice is not obvious — the fastest one depends on which phase your problem repeats.

Pick by the phase that repeats

A Newton loop re-factorizes every iteration → PARDISO (220× over SuperLU, vs cuDSS's 115×). A shift-invert eigensolve or constant-operator transient re-solves against one factorization → cuDSS, ~11× faster per solve and it takes a whole block of right-hand sides at once.

There is deliberately no auto: which wins depends on hardware jNO cannot inspect. Install both with pip install jax-numerical-operators[fem].

backend="host" — SuperLU on the host

factors on the HOST (SuperLU) and drives it from the device; same answer and same gradients (wrapped in custom_linear_solve, transpose via SuperLU trans="T"). Measured faster where cuSolver also works — Stokes 21,839 DOFs 0.27 s vs 1.67 s, H(curl) 17,072 complex DOFs 13.3 s vs 36.4 s — and it runs meshes cuSolver rejects (Stokes 26,908, H(curl) 26,154, both of which fail on GPU). Affordable because a direct solve factorises once: the operator crosses PCIe once, not per iteration. Read the win as cuSolver's sparse LU is weak, not GPUs lose — see the row below

backend="cudss" — NVIDIA cuDSS

NVIDIA cuDSS — the fastest direct backend wherever it runs, because it separates the symbolic plan from the numeric factorization and jNO caches on the sparsity, so the plan survives a change of values. Against backend="host" on an RTX 3070 (fp64 at 1/64 rate — the unfavourable card): Stokes saddle factorization 3.4 ms vs 79.9 ms, lap3d 50³ 576 ms vs 64,856 ms, and 64.7× per Newton step at n=64,000. Also factors the Stokes saddle cuSolver calls singular, with smaller residuals. Needs the optional stack (nvmath-python, cudss, cupy); raises a clear ImportError otherwise. Fill-in still governs 3-D (69×→218× nnz growth at lap3d 20³–40³), so it moves the ceiling and makes device memory the binding constraint — it is not a substitute for a preconditioner Two things happen automatically: a block right-hand side (n, k) is solved in one call (measured 2.7× at k=4 to 5.4× at k=16 over the same factorization solved column by column — this is what the shift-invert eigensolver's block Krylov sweep needs), and an exactly symmetric operator is factored as LDLᵀ instead of general LU (1.41× faster, 1.38× less peak device memory on lap3d 40³). Symmetry is tested to within a few ulps rather than bitwise — an assembled FEM tangent for a symmetric form is symmetric only up to assembly round-off (a vector element block contracts components in a different order for (a,i),(b,j) than for (b,j),(a,i), measured 0.25 ulps), so a bitwise rule sent every vector and coupled problem down the general-LU branch. Within the gate the two triangles are averaged, a correction bounded by the asymmetry it removes and therefore no larger than the round-off already in the matrix; outside it nothing is touched. The margin is not a knife-edge: the weakest genuine asymmetry that can be constructed — an advection coefficient of 1e-12, meaningless beside the Laplacian — is already 191 ulps. Measured end-to-end on a 3-D vector tangent, general → symmetric is 1.07–1.13× at 3.4k–11.8k DOFs (growing with size) and 1.10–1.62× on lap3d 20³–50³. The rest of the rule stands: symmetry is still tested and SPD is never inferred — a matrix symmetric only to ~1e-15 would otherwise be silently factored as (A+Aᵀ)/2, and a wrong SPD guess returns NaN. A replaced pivot is first given iterative refinement against the true matrix (Wilkinson 1963 — the remedy cuDSS's own docs prescribe), so a zero-pivot-but-nonsingular operator or a marginally conditioned saddle is solved (refined to ~1e-10 per solve, one SpMV per step) rather than refused; a genuinely singular operator — one whose refinement residual will not contract — still raises: cuDSS signals it through neither an exception nor a NaN, returning a finite plausible vector instead

backend="pardiso" — Intel MKL PARDISO

Intel MKL PARDISO, multithreaded CPU. The fastest factorization of the four, and like cuDSS it splits symbolic analysis from numeric factorization, so a Newton step reuses the analysis. On lap3d 50³ (n=125,000) against single-threaded SuperLU's 65,212 ms: factorization 298 ms, Newton re-factorization 296 ms — 220×, where cuDSS reaches 115×. Its adjoint is cheaper too — Aᵀx = b comes from the same factorization rather than a second one. An exactly symmetric operator uses LDLᵀ (1.9× on lap3d 50³, 13× on a saddle), which needs the upper triangle with an explicit diagonal — without that a saddle's constraint rows come back empty and PARDISO rejects the matrix. Like cuDSS it returns finite garbage on a singular operator, so jNO checks the perturbed-pivot count and raises. pip install jax-numerical-operators[pardiso], x86-64

A direct backend can return finite garbage on a singular operator

Neither cuDSS nor PARDISO signals singularity through an exception or a NaN — both return a finite, plausible-looking vector. jNO checks (the refinement residual for cuDSS, the perturbed-pivot count for PARDISO) and raises instead.

Host SuperLU (backend="host", and the default lu() / newton(direct=True) on the CPU, which is JAX's spsolve calling scipy) fails in three different ways, depending on pivoting luck:

• it aborts ("failed to factorize matrix at line 413 in file …/dpanel_bmod.c", after a burst of BLAS "On entry to DTRSV parameter number 6 had an illegal value" lines printed from C, which jNO cannot suppress). fem.solve turns this into a RuntimeError naming a singular / structurally rank-deficient operator, its size, SuperLU's own message and the same hints the cuDSS error gives. Measured on a three-field u–p–s saddle with continuous P1 stress (a known-incompatible space, singular by construction) at 3334 DOFs.
• splu (backend="host") reports Factor is exactly singular. That is named the same way, with the exact matrix size.
• JAX's spsolve answers an exactly singular matrix with NaN. The Newton verdict then reports a non-finite residual and points at a singular tangent, not at max_steps.

The same form at 2582 DOFs or fewer factors without complaint and returns a finite answer that is wrong in the null space. Nothing per-solve detects that, because detecting it would cost a residual product on every solve. Inside a Newton solve the convergence verdict catches it only if the wrong null-space component also stops the residual from falling. A consistent singular system can still converge, with an arbitrary null-space component. Outside Newton, check ‖Ax − b‖ yourself, or use cuDSS or PARDISO, which check. Naming the error costs nothing on a successful solve, because it runs only on the failure path.

Batched solves and operator storage — jno.setup(lu_stack=, matvec_format=)

Every lu backend has a vmap rule, so jax.jacrev / jax.jacfwd through a solve and a vmap over solves work on all four. host, cudss and pardiso factor once for a batch against one matrix and solve the whole block of right-hand sides in one call — the choice for Jacobians and sensitivities. The default device backend (cuSolver) takes one right-hand side, so it factors every system: one per call by default, or k per block-diagonal call with jno.setup(__file__, lu_stack=k) (or lu_stack = k under [jno] in .jno.toml). A larger k measured faster, but cuSolver allocates for the fill-in of all k systems, so the right k depends on the problem and the card; too large fails with a cuSolver allocation error, never a wrong answer.

The iterative solvers apply a sparse operator as cuSPARSE CSR or as split COO, chosen by timing both on the running device, once per operator size class (matvec_format="auto", the default). jno.setup(__file__, matvec_format="csr") or "coo" (or [jno] matvec_format) forces one — for reproducible timings, or to rule the choice out while debugging.

Preconditioners

Declarative specs — materialized against the assembled operator at solve time. A preconditioner never changes the converged solution, only the speed, so specs need no gradient path.

Spec	What it is	When
jno.precond.jacobi()	diagonal	the default — cheapest thing that helps
jno.precond.chebyshev(degree=…)	fixed-degree polynomial — matvecs and AXPYs only	GPU, where reductions cost; a fixed linear map, so it legally preconditions cg/minres
jno.precond.nystrom(rank=…)	randomized low-rank deflation	SPD only — a spectrum with a few large outlying eigenvalues, where jacobi can be worse than nothing
jno.precond.inner(solver)	any jno.solve solver as the M⁻¹ application	an inexact block/system solve — iterative inner ⇒ outer fgmres
jno.precond.form([...terms], inner=…)	an auxiliary operator assembled from ordinary traced terms	when you can write the preconditioner as a weak form
jno.precond.block_diag((field, spec), …)	per-field composition over fem.blocks (also triangular)	multifield problems
jno.precond.saddle(mass_weight=…, laplace_weight=…)	the standard saddle recipe as one call	Stokes / Biot — the common case
jno.precond.lsc() / saddle(schur="lsc")	least-squares commutator — the convection-aware Schur factor	Navier–Stokes, where the pressure-mass recipe degrades with Re
jno.precond.fsai(power=…)	factored sparse approximate inverse M⁻¹ = GᵀG — two sparse products, no triangular solves	SPD only — the rung between jacobi and multigrid, when the factor is reused (below)
jno.precond.schwarz(parts=…)	algebraic overlapping Schwarz (domain decomposition), two-level	SPD elliptic systems split across several devices; elasticity with nullspace="rigid" (below)
jno.precond.amg(cycles=…)	hybrid algebraic multigrid — host setup (pyamg), device apply	large SPD systems; needs the optional pyamg
jno.precond.jaxamg(symmetric=…)	GPU AMG via NVIDIA AmgX — setup and apply both on device	large systems on a GPU; needs the optional stack
.cached(refresh=…)	reuse an expensive setup across solves	Newton loops and transients, where the operator barely changes between solves
M1 + M2 / M1 @ M2	arithmetic on specs — applied in parallel and summed, or in sequence right-to-left	writing your own Schur approximation instead of asking for one

Every one of them is passed the same way — fem.solve(precond=…), alongside whichever linear= solver it is accelerating. Nothing else about the problem changes:

fem = jno.fem([ui.x * vi.x + ui.y * vi.y - 1.0 * vi, u(xb, yb) - 0.0])   # 379 DOFs, SPD

fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.jacobi())
fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.chebyshev(degree=4))
fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.nystrom(rank=20))
fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.amg(cycles=1))

# an ITERATIVE preconditioner needs a flexible outer solver — plain gmres assumes a fixed M
fem.solve(linear=jno.solve.fgmres(tol=1e-10),
          precond=jno.precond.inner(jno.solve.cg(tol=1e-2, maxiter=30)))

# reuse an expensive setup across a Newton loop or a march
fem.solve(linear=jno.solve.cg(tol=1e-10), precond=jno.precond.amg(cycles=1).cached())

# three more with the same shape, each needing something the others do not:
jno.precond.gmg(n_pre=2, n_post=2)      # geometric MG — needs a STRUCTURED grid
jno.precond.jaxamg(symmetric=True)      # GPU AMG (AmgX) — needs the [amg] extra and a CUDA device
jno.precond.ams()                       # H(curl) auxiliary-space MG — for N1E edge elements

All six agree — the preconditioner is a speed choice, never an answer

Solved on the same 379-DOF Poisson operator, every spec above lands within 3.5e-10 of the unpreconditioned reference (max|u − u_ref|). That is the contract: a preconditioner changes the iteration count, never the converged solution — which is also why these specs need no gradient path of their own.

A preconditioner's own inner solve is exempt from the convergence gate — inner(cg(tol=1e-2, maxiter=30)) asks for two digits deliberately. See the gate.

For a multifield system, compose per field over fem.blocks. The key is the trial symbol, not a name — the same u/p you wrote the weak form with:

fem.solve(linear=jno.solve.fgmres(tol=1e-10, restart=40),
          precond=jno.precond.block_diag(
              (u, jno.precond.inner(jno.solve.cg(tol=1e-2, maxiter=60))),   # velocity block
              (p, jno.precond.form([(1.0 / mu) * pp * qq], inner=jno.solve.dense())),  # pressure
          ))

Periodic ties, slip and hanging nodes. These make the solve run on the reduced system PᵀAP, reduced field by field. The block preconditioners split that operator: the slices they are handed (ctx.blocks) are each field's reduced DOFs, while fem.blocks keeps slicing the full solution that fem.solve returns. An auxiliary form([...]) is written on the full space as usual (no ties in it) and reduced with the same field's P. So block_diag, triangular, saddle (mass, Cahouet–Chabard, lsc), pcd and jno.solve.staggered are written exactly as on the untied problem. The preconditioners that build a pattern once (fsai, schwarz) build it from the reduced operator, and schwarz(nullspace="rigid") builds its rigid-body modes on the full mesh and restricts them to the kept DOFs (a translation along a periodic direction restricts exactly). One case is refused by name: a complex form solved as its fused real-equivalent [Re; Im] block. There a field is not one contiguous slice, so use a complex-native composition (a child such as ams()) or a whole-system preconditioner.

jno.precond.chebyshev(degree=…)

fixed-degree Chebyshev polynomial preconditioner: matvecs and AXPYs only, the GPU-era substitute for Gauss-Seidel/ILU smoothing, and a fixed linear map so it legally preconditions cg/minres. Spectrum bounds come from lmin/lmax when you pass them, else both ends are measured by Lanczos (Lanczos 1950, §II — the extreme Ritz values), at the same cost as a power iteration on the top end alone. This matters because the polynomial is a contraction only inside the interval it is fitted to: the lmin = lmax/30 guess, whenever the true ratio is smaller, leaves the lowest modes outside that interval where the polynomial amplifies them. Without the optional matfree package that guess is the fallback.

jno.precond.nystrom(rank=…)

randomized Nyström low-rank preconditioner (Frangella, Tropp & Udell, SIAM J. Matrix Anal. Appl. 44(2), 2023, Alg. 2.1 and §3), the rung between jacobi and a multilevel method. Sketches A against a random n × rank matrix — exactly rank matvecs, no factorization, no triangular sweep — and deflates the captured top of the spectrum. That is what jacobi cannot do: a diagonal rescales, it cannot separate a few large outlying eigenvalues, and on such a spectrum Jacobi can be worse than no preconditioner at all (measured on a rank-15-outlier SPD system: 124 CG iterations for jacobi, 98 unpreconditioned, 46 for nystrom(rank=20)). SPD only — the sketch takes a Cholesky, so an indefinite operator gives NaN rather than a quiet wrong answer.

The momentum block: Picard, not Newton

A block preconditioner needs an approximate F⁻¹ on the momentum block, and which linearisation you pick decides whether one exists. The full Newton tangent carries the reaction term (δu·∇)u, which is indefinite — measured on a lid-driven cavity, the smallest eigenvalue of its symmetric part goes negative by Re = 100:

Re = 100	min eig of sym(F)	AMG on that block
Newton tangent	−6.4e-04	rel-resid 1.2e-03, diverges by Re = 400
Picard / Oseen (jno.lag on the convecting velocity)	+8.5e-04	6.5e-11

Smoothed aggregation assumes a positive diagonal, so on the Newton tangent it builds a degenerate hierarchy — jno.precond.amg() now refuses that by name rather than returning a V-cycle that silently does nothing. Lag the convecting velocity and the block is definite and AMG-friendly. This is why the Navier–Stokes preconditioning literature builds on Picard.

Does it scale?

3-D Stokes, lid-driven cube, lu(backend="host") against fgmres + triangular((u, amg()), (p, lsc())), one 8 GB card:

DOFs	direct	iterative
15,468	2.1 s	2.9 s	direct still wins
29,114	11.5 s	4.6 s	2.5×
49,072	53.3 s	6.5 s	8.2×
76,542	166.4 s	11.4 s	14.6×
112,724	431.3 s	52.0 s	8.3×

The crossover is around 15–20k DOFs. Note what the direct solve is not doing: it does not run out of memory — it scales as roughly O(N³) in time and simply loses. Below the crossover it is the right choice and stays the default.

AMG on a march

jno.precond.amg() builds its hierarchy on the host from a concrete matrix, and a march linearises inside lax.scan where every operator is traced — so it used to refuse every transient problem. A march now builds it once, before the scan, from the step tangent (a state-dependent mass c(u) u_t included).

A frozen setup stays correct — a preconditioner changes convergence speed, never the answer — but goes stale when the operator moves. cached(spec, refresh=k) spells the cadence, and the march then runs as ceil(n/k) scans, re-composed against the carried state:

fem.solve(linear=jno.solve.fgmres(), precond=jno.precond.amg())                      # frozen once
fem.solve(linear=jno.solve.fgmres(), precond=jno.precond.cached(jno.precond.amg(), refresh=25))

Measured where a melt pool's coefficients move ~1e13 between solid and molten: frozen once, the Krylov residual stalls at 2.9e-03; at refresh=25 it reaches 1.5e-04. theta()/exponential() integrate the march themselves and are not chunked, so a cadence beside one raises rather than never firing.

A frozen preconditioner is probed once and refused if it makes a random residual worse — but only when it stands alone, since a block preconditioner need not reduce a full residual in one application.

A steady nonlinear solve has the same problem one level down: Newton runs as a lax.while_loop, so the tangent is traced there too. An unbuilt amg() (or ilu(), or a block holding one) is therefore built once, before the loop, from the assembled tangent at the initial guess x0 (zero by default; a periodic, slip or hanging-node system uses its reduced tangent). It does not follow the tangent as Newton moves: the root is unchanged and only the Krylov speed depends on how far the tangent drifts from x0. Specs that already materialize inside the loop (jacobi, form, lsc, pcd) are left as they were.

fem.solve(linear=jno.solve.fgmres(), precond=jno.precond.amg())                  # nonlinear: built at x0
fem.solve(linear=jno.solve.fgmres(),
          precond=jno.precond.triangular((u, jno.precond.amg()), (p, jno.precond.form([pb * qb / mu]))))

For Navier–Stokes, lag the convecting velocity (jno.lag, above) so the velocity block AMG is built on is the Oseen operator. cached(spec, refresh=k) is refused here rather than ignored: a rebuild every k Newton iterations would need the loop cut into chunks, which the steady solve does not do. To rebuild at a better state, solve again from that state (x0=).

When the pressure mass is not enough

saddle()'s pressure-mass Schur approximation stands in for the Schur complement of a viscous operator. Once convection dominates it degrades. Measured on the Newton tangent of a lid-driven cavity (812 DOFs, momentum block solved exactly so the count isolates the Schur factor):

Re	10	50	100	400	1000
saddle() pressure mass	26	38	82	212	105
lsc()	112	108	106	104	284

Over Re = 10 → 400 the mass degrades 8× while LSC is flat. The robustness ends: by Re = 1000 LSC has jumped to 284 and the mass has (non-monotonically) improved, so LSC is no longer ahead. This is a result about a range, not an unconditional one. It is built from the system's own blocks — S⁻¹ ≈ (BM⁻¹Bᵀ)⁻¹(BM⁻¹FM⁻¹Bᵀ)(BM⁻¹Bᵀ)⁻¹ — so unlike the form()-based approximations it cannot be written by the user as weak forms, which is why it ships as a built-in. B and M do not depend on the solution, so P = BM⁻¹Bᵀ is assembled sparsely and factorised once, for every application and every Newton step; only F varies, and it is only ever applied.

Measuring this at all needs args=

These counts come from fem._op.jacobian(sol, args={"nu": …}). If the viscosity is a runtime parameter and you call jacobian(sol) without args, it evaluates at an unset parameter and returns a matrix that is 25 % wrong — silently. The operator then does not change with Reynolds number at all, and every preconditioner looks perfectly robust on it. An earlier version of this table was measured that way and was wrong.

What LSC does not give you

It is not mesh-independent here. At fixed Re = 100 on a cavity, iterations grow 36 → 67 as the pressure space grows 30 → 198 (the pressure mass grows 32 → 91 over the same range, so LSC is better but neither is flat). And the Reynolds result above is resolution-dependent: at 613 DOFs the same sweep reverses the ordering, because that mesh never enters the regime where the mass approximation degrades. Read the table as a resolved-discretisation result, not an unconditional one.

PCD — and the boundary condition that decides it

PCD (Kay, Loghin & Wathen 2002) is the other convection-aware Schur approximation, S⁻¹ ≈ A_p⁻¹ F_p M_p⁻¹. It ships as jno.precond.pcd():

xin, yin, _ = d.variable("inlet", split=True)
fem.solve(linear=jno.solve.fgmres(),
          precond=jno.precond.triangular((u, jno.precond.amg()),
                                         (p, jno.precond.pcd(viscosity=nu, inflow=(xin, yin)))))

inflow is the coordinate tuple, not a region name — domain.variable(...) mints a fresh region on every call (inlet_3, inlet_4, …), so a name re-resolved inside the spec would point at a different region than the one your form constrained and the Dirichlet row would quietly miss the boundary. Passing a string is refused by name.

Unlike LSC every factor is a weak form, so the supplied spec is a convenience, not a capability — it is verified equal (to 7e-16) to the same thing written by hand, which is where to go for a variant it does not cover:

w  = wv.bind(x=xi, y=yi).freeze(u_lagged)          # the convecting field, as KNOWN nodal data
lap_p, conv_p = pp.x*qq.x + pp.y*qq.y, (w[0]*pp.x + w[1]*pp.y)*qq

Ap = jno.precond.form([lap_p,              p(x_in, y_in) - 0.0], inner=host)   # A_p⁻¹
Fp = jno.precond.form([nu*lap_p + conv_p,  p(x_in, y_in) - 0.0], inner=False)  # F_p, APPLIED
Mp = jno.precond.form([pp*qq],                                    inner=host)   # M_p⁻¹
schur = Ap @ Fp @ Mp

Two conditions decide whether it works at all, and both are easy to get wrong:

Inflow/outflow problem	On an enclosed flow (a lid-driven cavity) PCD fails outright — 400 iterations, the cap, at every Reynolds number, and its eigenvalues straddle zero. The commutator argument leans on boundary behaviour an enclosed flow does not provide. Use lsc() there.
Inflow-Dirichlet on A_p and F_p	With Neumann pressure operators PCD also fails (400). With a Dirichlet condition on the inflow boundary it takes 77. Same problem, same mesh, same everything else.

Channel flow, momentum block exact, per Reynolds number:

Re	10	50	100	400	1000
pressure mass	41	184	301	284	318
lsc()	63	158	150	291	263
PCD	37	75	77	237	355

PCD is decisively best from Re = 10 to 100 — 4× the mass at Re = 100 — and loses its edge by Re = 400, where everything degrades. Neither method is Reynolds-robust at the top of that range.

Preconditioners compose like the operators they stand for

A materialized spec is a linear map v ↦ M⁻¹v, so the classical physics-based Schur approximations are arithmetic on specs rather than a menu of built-in names:

Ap = jno.precond.form([inner(gp, gq, n_contract=2), p_sym.pin()])   # A_p⁻¹  (gauged Laplacian)
Mp = jno.precond.form([pp * qq])                                    # M_p⁻¹  (pressure mass)
Fp = jno.precond.form(F_terms, inner=False)                         # F_p    APPLIED, not inverted

schur = Mp + Ap        # Cahouet & Chabard (1988) — which `saddle(laplace_weight=…)` already is
schur = Ap @ Fp @ Mp   # PCD — Kay, Loghin & Wathen, SIAM J. Sci. Comput. 24 (2002) 237

fem.solve(linear=jno.solve.fgmres(),
          precond=jno.precond.triangular((u, jno.precond.amg()), (p, schur)))

inner=False is the piece that makes a product expressible: it applies the assembled form (M v = A v) instead of inverting it, which is what the middle factor of PCD needs. On its own such a factor is not a preconditioner and will make a solve worse — it exists to be composed.

A product is not symmetric

(AB)ᵀ = BᵀAᵀ, and a product of two SPD inverses is not SPD. Drive a product with jno.solve.gmres() / fgmres(), never cg/minres — a short-recurrence method will simply fail to reach the residual gate, loudly. (The transpose applier reverses the factor order for you, which is what keeps the reverse pass of a differentiable solve preconditioned.)

This is why jno.precond.saddle() has no schur= argument: the approximations that would populate such a menu are things you can write.

jno.precond.form([...terms], inner=…)

preconditioners as weak forms: assemble an auxiliary operator from ordinary traced terms and invert it as M⁻¹ (weighted mass matrices, shifted-Laplacian Helmholtz twins, low-order proxies — written in the PDE's language). Pass a callable form(lambda sol: [...terms...]) for a solution-dependent auxiliary — a (1/μ(u))-weighted Schur mass whose weight is computed from the current solution. It is re-assembled once per outer solve from that solve's entry iterate (warm start / previous march step) — the Picard-lagged preconditioner (Elman, Silvester & Wathen 2014, §9.2): the coefficient trails the solution by one outer solve, which changes convergence speed only, never the answer. Eager by necessity — every Newton loop is a lax.while_loop, so the per-step iterate is a tracer no host assembly can see; a fully traced solve that never supplies a concrete iterate gets a loud NotImplementedError, not a garbage preconditioner.

A coefficient from another field. An auxiliary form is written over one field, and the coefficient it needs often lives on another — PCD's F_p carries the velocity into a pressure form, a lagged eddy viscosity carries the previous velocity, a wall-distance field is reused across forms. Freeze the foreign field and use it like any other coefficient:

w_frozen = w.bind(x=xi, y=yi).freeze(w_values)          # another form's solved field, as data
Fp = jno.precond.form([(nu + w_frozen) * (pi.x * qi.x + pi.y * qi.y), p(xb, yb) - 0.0])

The auxiliary form never declares that field as an unknown, so it is resolved by space rather than by identity: a field of the same space and order has the same nodes and connectivity on this mesh, so the gather is exact, not an interpolation.

Matching space only — a mismatch is refused, not approximated

The foreign field's space must be one the auxiliary form already uses. A P2 field in a P1 form raises: P2 values carry edge nodes a P1 gather never indexes, and reading them through a P1 slot would produce a plausible, wrong operator rather than a visible failure. Project onto the form's own space first, or write the form over a field of the matching space.

This is what still blocks PCD on a Taylor–Hood pair, where the velocity is P2 and the pressure P1. lsc() needs no such coefficient, which is why it exists.

jno.precond.block_diag((field, spec), …)

/ jno.precond.triangular((field, spec), …) — per-field composition over fem.blocks. triangular is the standard saddle-point shape: last block solved first, substituted back through the assembled off-diagonal matvecs.

jno.precond.saddle(mass_weight=…, laplace_weight=…)

that standard shape as one call, for the common case. It finds the constraint block structurally (the field with no diagonal entry — the same detection behind the saddle-system warning), puts amg() on the momentum block and a weighted pressure mass matrix on the constraint block, and assembles that mass on the domain's own P1 space, so no symbols are passed and it reads identically in 2-D and 3-D. mass_weight is explicit and never inferred (1/μ for Stokes): digging μ out of an arbitrary weak form is fragile, and a variable viscosity would silently take the wrong weight — which costs iterations without ever failing. A wrong weight changes convergence speed only, never the answer. Pair with jno.solve.fgmres; minres does not apply (block-upper-triangular is nonsymmetric). Needs pyamg, and refuses by name on a non-saddle system or a constraint field that is not P1 — for anything outside that, compose triangular yourself. The mass matrix alone is the pure-Stokes approximation: a strong reaction term (Brinkman / Darcy drag, or the 1/dt mass of a small implicit step) makes the Schur complement stop looking like a mass matrix, and the mesh-robustness degrades. laplace_weight switches to the Cahouet–Chabard approximation, a pressure mass plus a pressure Laplacian, S⁻¹ ≈ μ·M_p⁻¹ + α·L_p⁻¹ (Cahouet & Chabard, IJNMF 8, 869–895, 1988, §3). Both weights are the reciprocal of the coefficient the term stands for, so mass_weight=1/μ and laplace_weight=1/α. Measured through this spec on a 2-D Brinkman channel (μ=1, mesh_size=0.12, preconditioned GMRES to 1e-8, restart=200): α = 0 / 1e2 / 1e3 / 1e4 → 83 / 135 / 503 / 2312 with the mass alone, 83 / 60 / 76 / 74 with Cahouet–Chabard. The point is the flatness, not the 31× at α=1e4 — the count stops tracking α. The same collapse appears with the momentum block inverted exactly instead of by AMG (31 / 63 / 134 / 189 → 31 / 26 / 25 / 24), which is how you can tell it is the Schur approximation doing the work and not the multigrid. At α=0 the two are the same preconditioner — that column is one number measured twice — so leave laplace_weight at its None default on pure Stokes and the Laplacian is never assembled. The pressure Laplacian is pure-Neumann and therefore singular — and a sparse LU of it factors happily and then applies nonsense, so the auxiliary carries its own gauge pin and the applier projects the constant out on both sides. It is derived for a constant α; a spatially varying drag (topology optimisation's α(ρ)) takes a representative scalar and degrades gracefully rather than failing.

jno.precond.amg(cycles=…)

hybrid algebraic multigrid: setup once on the host via the optional pyamg (Vaněk, Mandel & Brezina, Computing 56, 1996; PyAMG — Bell et al., JOSS 8(87), 2023), applied as a pure-JAX V-cycle with Chebyshev smoothing (Adams et al., JCP 188, 2003). The apply is jit/vmap-native and exactly linear, so it preconditions cg/minres too. Mesh-independent convergence ⇒ the choice for large elliptic blocks. Inside traced/parametric solves, pre-build eagerly: spec = jno.precond.amg(); spec.build(fem.A). Without pyamg, amg raises a clear install hint.

jno.precond.jaxamg(symmetric=…)

GPU AMG via NVIDIA AmgX as the M⁻¹ application (setup and apply both on the device). symmetric=False builds a second hierarchy on Aᵀ so the adjoint (reverse-mode) solve of a non-symmetric operator is preconditioned too. Measured caveat: each AmgX application carries ~11 ms of fixed handle overhead, so it pays only where the iteration savings exceed it — for most problems prefer precond.amg() (whose V-cycle compiles into the solve) or linear=jno.solve.amg() (ONE AmgX crossing per solve, with warm structure-keyed re-setup measured ~10x cheaper on repeats).

.cached(refresh=…)

on any spec — reuse an expensive setup across solves: False frozen, True rebuild on shape/sparsity change, an int k to rebuild every k-th materialization (the cadence for a march whose operator values drift step by step), or a ctx -> key callable.

The flagship pattern — Taylor–Hood Stokes by FGMRES with an inexact velocity block solve and the viscosity-weighted pressure-mass Schur approximation (Elman, Silvester & Wathen, 2014, §9.2):

sol = fem.solve(
  linear  = jno.solve.fgmres(tol=1e-10, restart=40),
  precond = jno.precond.triangular(
      (u, jno.precond.inner(jno.solve.cg(tol=1e-2, maxiter=60))),   # Â⁻¹: inexact CG
      (p, jno.precond.form([(1.0/mu) * pp * qq], inner=jno.solve.dense())),  # Ŝ ≈ μ⁻¹ M_p
  ),
)

…and the same recipe when the defaults suit — multigrid on the momentum block, exact pressure mass:

sol = fem.solve(linear=jno.solve.fgmres(tol=1e-10, restart=150),
              precond=jno.precond.saddle(mass_weight=1.0/mu))

Give FGMRES enough restart: the default restart=30 stagnates on a 3-D Taylor–Hood block preconditioner and does so quietly — it still returns, just slower and less accurate. Measured at 4302 dofs, tol=1e-10: restart=30 → 2.30 s for 3.3e-6; restart=150 → 0.49 s for 3.3e-9. On 3-D Stokes this is what makes the recipe overtake a direct factorisation — measured crossover at ~15k dofs, and 3.4–3.9× faster (23.2 s → 5.9–6.9 s over repeat runs) at 41k, where LU's fill-in also costs more memory.

On a reaction-dominated system — Brinkman/Darcy drag, or a small implicit time step — add the Laplacian leg, and the iteration count stops tracking the reaction coefficient:

sol = fem.solve(linear=jno.solve.fgmres(tol=1e-10, restart=150),
              precond=jno.precond.saddle(mass_weight=1.0/mu, laplace_weight=1.0/alpha))

Single precision inside a double solve — float32=True

Every jno.precond constructor takes float32=False. With True, the preconditioner is built against a single-precision copy of the operator and applied in single precision, while the Krylov solve — its residual, its dot products and its convergence test — stays in double. The answer meets the same tolerance; only the preconditioner's arithmetic is cheaper.

fem.solve(linear=jno.solve.cg(tol=1e-8), precond=jno.precond.amg(float32=True))

A single-precision preconditioner is only approximately symmetric, and standard preconditioned CG relies on exact symmetry: it lost conjugacy and stopped at a true residual of 6e-6 against a requested 1e-8 (float32 FSAI, 3-D elasticity), silently. jno.solve.cg therefore switches to flexible CG (Notay, SIAM J. Sci. Comput. 22(4), 2000) whenever the preconditioner is float32 (or non-symmetric), which keeps the requested accuracy.

CG to 1e-8, preconditioner in float32 vs float64 (RTX 3070)

Measured on 3-D P1 (303k), P2 (142k) and elasticity (206k) operators:

• amg: 1.49×, 1.55×, 1.2–1.3×
• chebyshev: 1.8–3.0×
• fsai: 1.1–1.2×
• jacobi: 1.0× — a single scale, nothing to save

The gain is memory bandwidth: it is largest where the preconditioner dominates the iteration. One device, one precision pair — measure on yours.

FSAI — jno.precond.fsai()

A factored sparse approximate inverse for SPD operators: M⁻¹ = GᵀG ≈ A⁻¹, with G lower triangular on the pattern of tril(A) (or tril(A²) with power=2) and each row a small independent dense solve (Kolotilina & Yeremin, SIAM J. Matrix Anal. Appl. 14(1), 1993). Applying it is two sparse products — no triangular solves and no sequential sweep, unlike ILU — so it runs well on a GPU.

The pattern is computed once on the host; the numeric factor follows the operator's current values inside a compiled solve, a Newton loop or a march. A non-symmetric operator is refused, and a changed sparsity pattern poisons the factor with NaN rather than applying a stale one.

CG to 1e-8, RTX 3070, float64

	3-D P1, 303k	3-D P2, 142k	3-D elasticity, 206k
jacobi	216 its / 76 ms	197 its / 58 ms	778 its / 588 ms
fsai	89 its / 58 ms	82 its / 39 ms	325 its / 387 ms
amg	48 ms	52 ms

The setup costs about one to two solves (101, 115, 530 ms there), so FSAI pays when the factor is reused: a constant-operator march, .cached() across solves, several right-hand sides. power=2 grows the setup steeply (7.1 s on the P2 case, for 32 iterations) — start with power=1.

Overlapping Schwarz — jno.precond.schwarz()

Domain decomposition as a preconditioner, built from the assembled operator alone: no mesh, no physics, so any FEM or FDM system qualifies. (This is not the multi-method coupling of Domain decomposition, which couples different solvers on different regions.)

1. Partition. METIS splits the operator's graph into parts compact pieces with a minimal cut (default: one per ~256 unknowns). Unknowns with no neighbours — eliminated Dirichlet rows — are left out and solved by their diagonal.
2. Overlap. Each piece grows by overlap layers of graph neighbours (default 1).
3. Local solves. Each subdomain's matrix is A restricted to its unknowns (a zero-Dirichlet problem on its grown boundary). All of them are solved exactly and together, as one batched product with their inverses. The results are added back, summing where subdomains overlap.
4. Coarse correction (coarse=True, the default). One small global problem with a few unknowns per subdomain carries information across the whole domain in one step. Without it, the iteration count grows with the number of parts; with it, it stays flat.

fem.solve(linear=jno.solve.cg(tol=1e-8), precond=jno.precond.schwarz())
fem.solve(linear=jno.solve.cg(tol=1e-8), precond=jno.precond.schwarz(nullspace="rigid"))  # elasticity
fem.solve(linear=jno.solve.cg(tol=1e-8), precond=jno.precond.schwarz(), shard=4)          # 4 devices

On jno.fdm it composes the same way, on the operator FDM assembles for the slots (FDM solver slots), including inside a jno.core training step and for vector unknowns with nullspace="rigid":

jno.fdm([...]).solve(linear=jno.solve.cg(), precond=jno.precond.schwarz())

Argument	Meaning
parts	number of subdomains; default ~256 unknowns each, rounded up to a multiple of the device count
overlap	layers of neighbours each subdomain grows by
coarse	two-level (default) or one-level
restricted	RAS (Cai & Sarkis 1999): usually fewer iterations but non-symmetric — cg then runs flexible CG
nullspace	what the coarse space carries per subdomain: None one constant (Nicolaides 1987); "rigid" the rigid-body modes of a displacement field (translations + rotations), which elasticity needs; or an explicit (n, k) array
float32	single-precision application

CG to 1e-8, two-level, RTX 3070 — iterations at 64 / 256 parts

Problem	Iterations	Solve
2-D Poisson, 11.8k	50 / 43	32 / 24 ms
2-D cantilever, nullspace="rigid"	57 / 41	44 / 55 ms
3-D Poisson, 18.9k	37 / 34	130 / 63 ms

On the cantilever the coarse space decides the scaling: at 16 / 64 / 256 parts, one constant per part needs 154 / 207 / 256 iterations, the rigid-body modes 103 / 90 / 56.

On one GPU it is slower than amg (2-D Poisson 11.8k, 256 parts: 43 its / 24 ms against amg's 12 its / 3.6 ms). Its value is placement across devices.

Several devices (fem.solve(shard=…), a concrete or a traced operator): the subdomains are split into one block per device. Each device builds, inverts and applies only its own, and their contributions are combined by an all-reduce of an n-vector (three per application in the two-level form). The coarse matrix and the vectors stay replicated. Two limits:

• The numeric setup gathers the operator's pattern and values onto every device, once per factorisation (O(nnz) there, transiently). The O(parts · m²) local matrices are what gets split.
• Verified on simulated CPU devices (per-device block counts in the compiled program, the all-reduce, answers and a gradient through the traced route on 1, 2 and 4 devices). Not timed on multi-GPU hardware.

Scope. Schwarz suits elliptic, positive-definite systems: diffusion, elasticity, implicit time steps. Indefinite problems (Helmholtz, a saddle point taken as a whole) are outside what it handles robustly. It needs x64 and an assembled operator, and METIS: pip install "jax-numerical-operators[fem]" (or [fdm], or [metis]).

References: Toselli & Widlund, Domain Decomposition Methods, Springer 2005; Karypis & Kumar, SIAM J. Sci. Comput. 20(1), 1998 (METIS); Vaněk, Mandel & Brezina, Computing 56, 1996 (rigid-body near-null space).

Nonlinear drivers

The default Newton — jno.solve.newton()

A nonlinear fem.solve(), and each step of a nonlinear march, runs jno.solve.newton() unless told otherwise. Its direct= argument picks how each linear step is solved:

direct=	tangent	linear step
None (default)	assembled J(u) (a BCOO), wherever the assembler provides one	Jacobi-BiCGStab, or the linear=/precond= slots. The BiCGStab answer is verified (one SpMV) and re-solved with Jacobi-GMRES when its relative residual is not small
False	matrix-free — J·v from a JVP, no tangent stored	BiCGStab, or the linear= slot
True	assembled	a sparse LU each step, or the linear= slot's direct solver (see Sparse-direct Newton below)

When each one falls back:

• The default runs matrix-free, as direct=False does, when the problem has no assembled tangent (a residual-only operator — the matrix-free-only paths listed under the sparse-direct Newton below).
• With nonlinear= unset, a direct linear= slot (lu, dense, amg) selects direct=True.
• direct=True needs the assembled tangent and raises without one; an explicit matrix-free nonlinear= beside a direct linear= raises too.

Because the default hands its linear step an assembled matrix, a precond= that needs one (jacobi) composes with it; direct=False has no matrix and is the choice when the assembled tangent would not fit in memory. All three modes are differentiable (implicit differentiation through lax.custom_root).

Picard / lagged coefficients — jno.lag

When a solution-dependent coefficient's Newton tangent destroys the linearized system's structure (the classic case: a shear-thinning viscosity μ_eff(u) in non-Newtonian Stokes flow, whose full-Newton velocity block defeats AMG/block preconditioners), freeze it with jno.lag(...) and drive with jno.solve.picard():

fem = jno.fem([2 * jno.lag(mu_eff) * inner(eps(ui), eps(vi)) - pi * div(vi), ...])
sol = fem.solve(nonlinear=jno.solve.picard(damping=0.7), linear=jno.solve.fgmres(),
                precond=jno.precond.triangular(...))

lag is stop_gradient on the traced expression, so the residual's linearization is the Picard operator — each outer step re-solves the lagged system (linear convergence, but every inner system keeps its symmetry/definiteness); the converged solution is identical to full Newton's. Without any lag marker, picard(damping=…) is exactly damped Newton (jno.solve.newton(damping=…)). Caveat for inverse problems: implicit differentiation then also uses the lagged Jacobian — drop lag when exact parameter gradients matter more than per-step solvability.

Anderson acceleration — picard(anderson=m), staggered(anderson=m)

A fixed-point loop converges only linearly, and slowly when the lagged coupling is strong. anderson=m mixes the last m iterates into each new one (type-II Anderson acceleration, Walker & Ni, SIAM J. Numer. Anal. 49(4), 2011). It is safeguarded: a mixed step that does not decrease the residual retreats toward the plain iterate, so it is never worse than anderson=0, the default. Gradients are unchanged, since the converged point is still differentiated through custom_root.

sol = fem.solve(nonlinear=jno.solve.picard(anderson=5))
sol = fem.solve(nonlinear=jno.solve.staggered([u, d], anderson=5))

-∇·((1 + c u²)∇u) = f, lagged coefficient, Picard steps to rtol 1e-8

c	picard()	picard(anderson=5)
5	33	11
20	fails in 400 steps	14
50	fails in 400 steps	20

Alternate minimization — jno.solve.staggered([u, d]) / staggered([[v, p], [T]])

Some coupled energies are non-convex in the fields jointly but convex in each separately. A monolithic Newton then has no descent guarantee and simply diverges; solving one field at a time turns the problem into a sequence of convex solves:

sol = fem.solve(nonlinear=jno.solve.staggered([u, dm]))   # sweep u, then dm, until both converge

Variational phase-field fracture is the canonical case — (1-d)²|∇u|² is quartic in the pair, while u alone is a linear elasticity problem and d alone a linear elliptic one. Measured on the coupled damage form in tests/test_fem_staggered.py: newton() leaves with a residual around 1e+25 (it raises), and the staggered sweep converges to a genuine root of the coupled system. Fixed-stress Biot poroelasticity and thermo-mechanical staggering have the same shape.

Algorithm: Bourdin, Francfort & Marigo, Numerical experiments in revisited brittle fracture, JMPS 48 (2000), §3 — as the staggered operator split with a history field, Miehe, Welschinger & Hofacker, IJNME 83 (2010).

Groups — fields solved together inside one sweep. An entry of the field list may itself be a list:

fem.solve(nonlinear=jno.solve.staggered([[v, p], [T]], direct=True), linear=jno.solve.lu())

This is not a convenience. A velocity/pressure pair cannot be swept apart: the pressure block is the constraint block, with no diagonal of its own — the block jno.precond.saddle locates structurally — so "solve p with v frozen" is not a well-posed sub-problem. Any flow staggered against a solid or a temperature therefore has to group its Stokes pair. Measured on a three-field Stokes/temperature problem (tests/test_fem_staggered_groups.py): the grouped sweep lands on the monolithic root to 1e-7 relative, and the flat staggered([v, p, T]) does not converge at all.

A bare symbol among lists is its own group, so [[v, p], T] is legal and [u, dm] keeps meaning exactly what it always did — one field per sweep. Every block must still appear, in exactly one group.

A group holding a constraint field is indefinite, so it wants direct=True for the reason in the next paragraph: a matrix-free sub-solve cannot be preconditioned, and unpreconditioned BiCGStab on a saddle block is not a solve. That is documented rather than enforced — fem.solve's own saddle warning already fires on the shape.

direct=True factorizes each field's diagonal block rather than solving it matrix-free, and pairs with a linear= slot:

fem.solve(nonlinear=jno.solve.staggered([u, dm], direct=True), linear=jno.solve.lu(backend="pardiso"))

It exists because the matrix-free sub-solve cannot be preconditioned. A precond= spec materializes against an assembled operator, and a staggered sub-problem is a restriction closure (x -> R(u with block set to x)[block]) with no matrix, so an ill-conditioned block is solved by unpreconditioned BiCGStab — near-incompressible elasticity (ν → 0.5) being the usual victim. That was also why a direct linear= slot used to be refused against staggered outright: there was nothing to factorize.

The extraction avoids a data-dependent nnz (which would break the static shapes a traced Newton needs): instead of slicing J[b][:, b] out, the complement's rows and columns are zeroed and given a unit diagonal, so the block is solved as [[J_bb, 0], [0, I]] against [-r_b, 0]. The padding is pure diagonal — no fill-in — so the factorization cost stays the block's. With bounds, the matrix handed to the factorization is the min-map's semismooth Jacobian (identity rows on the active set), which jax.linearize derives for free on the matrix-free path but an assembled tangent does not.

3-D Yeoh phase-field march — 576 DOFs, 8 load steps, CPU

Measured on a 3-D Yeoh phase-field march (576 DOFs, 8 load steps, CPU): 42.80 s → 10.78 s, 4.0×, to the same answer. Two honest caveats: the full tangent is assembled to use one block of it (a sparsity-caching backend — pardiso/cudss — then pays only the numeric re-factorization per sweep), and on a well-conditioned problem the matrix-free default is cheaper. This is not a free upgrade and it is not the default. The load-path march threads the tangent too, which it previously never did: before this, a direct driver inside a tau= path silently fell back to the matrix-free inner solve.

Line search

line_search= globalizes each sub-solve's Newton steps, and is on by default:

value	what it does
"backtrack" (default)	residual-norm Armijo, halving from damping. What this used to do under line_search=True.
True	exact line search — bisect for the root of the directional derivative R(x+λd)·d, i.e. the minimizer of the energy along the Newton direction
False	no line search; take damping

The default moved from False to "backtrack" because no line search is a genuine footgun on finite-strain forms: measured here, a 3-D Yeoh P2 march produced NaN on load step 1 without one (the undamped step inverts an element, det F ≤ 0, and J^(-2/3) is NaN) and solved cleanly with one.

True implements Heinzmann, Vicentini, Carrara et al., Iterative convergence in phase-field brittle fracture computations: exact line search is all you need, Computational Mechanics (2026), arXiv:2511.23064, §3 Algorithm 2 — the same algorithm they contributed to PETSc as SNESLineSearchBisection. Their Props. 1–2 and Remark 4 chain into a convergence guarantee for the whole alternate-minimization scheme, provided each sub-problem is strictly convex and coercive (jNO's bounds min-map is semismooth, so that proof does not cover it — the same gap they note for reduced-space active sets).

It is not the default, because we have not measured it beating backtracking. On the problems tested here — a 2-D SENT plate with and without a volumetric-deviatoric split, at several load levels — both need the same 21–22 staggered sweeps, and the exact search costs ~15% more wall time from the extra residual evaluations. Note the paper's own failure cases arise only where the mechanical sub-problem is non-linear (their §: the residual "reduces to an affine form in the absence of an energy decomposition"), at critical load steps reached along a path — a regime not reproduced here. Reach for True when a sub-solve stalls; the theory is on its side even where our measurements are neutral.

Over-relaxation

Alternate minimization is a nonlinear block Gauss–Seidel iteration, so over-relaxation accelerates it exactly as it does the linear one — go ω times as far along each sub-step's own update direction, per block:

fem.solve(nonlinear=jno.solve.staggered([u, dm], over_relax=1.4))

Algorithm: Farrell & Maurini, Linear and nonlinear solvers for variational phase-field models of brittle fracture, IJNME 109 (2017) 648–667, Algorithm 2 (ORAM), §2.1. ω = 1 is plain alternate minimization. Kahan's classical bound gives ω ∈ (0, 2) as necessary for SOR to converge, and anything outside raises.

Whether it pays is problem-dependent, and there is no way to know in advance. The paper is explicit — they "rely on the naïve strategy of numerical experimentation on coarser problems" and defer automatic selection to future work. Their own results split cleanly: Table I (a propagating crack, where AM converges slowly) gains 58–73% fewer iterations; Table II (where AM already converges fast) gets 0% — over-relaxation hinders it, going 37 → 111 → 185 → 326 → 747 iterations as ω runs 1.0 → 1.8. On a 2-D small-strain phase-field problem that drives damage to 1, ω = 1.4 was 1.95× faster (3096 → 1587 ms, warm median of 3). So it defaults to 1, and a short ω sweep on a coarse version of the problem is the only way to know — the paper's own two tables disagree with each other.

ω cannot diverge. The extrapolation is not taken on trust: the step retreats by bisection on [1, ω] until the trial point has a finite residual and is feasible, and ω = 1 is the sub-solve's own converged answer, already evaluated and therefore admissible by construction. So the worst case is that the sweep degrades to plain alternate minimization. This matters most on finite-strain forms, where an unguarded step past a converged answer inverts an element (det F ≤ 0, so J^(-2/3) is NaN) — measured on the 3-D Yeoh SENT march, which NaN'd on load step 1 at every ω down to 1.1 before the guard and runs to completion with it.

The test is finiteness and feasibility, not descent. Over-relaxation is deliberately not a descent method, so demanding a residual decrease would reject almost every ω > 1 and silently collapse the feature back to ω = 1 while appearing to keep it. Feasibility is the paper's own rule (§2.1 backs ω off by bisection on [1, ω] for the bound constraint); finiteness is the generalization.

Over-relaxation also acts on the free DOFs only. Farrell & Maurini's ũ lives in the constrained space C_ū, where a prescribed DOF has δ = 0; jNO imposes essential conditions as residual rows, so without the mask the sub-solve's exact hit on the prescribed value gets extrapolated past — measured on one row with g = 2, ω = 1.7 gave 3.40 → 1.02 → 2.69, an oscillation decaying only as |1−ω|ᵏ, worst on a ramped condition where g moves every load step. The mask covers every nonlinear problem — plain or parametric, 2-D/3-D or 1-D (non-nodal families excepted) — and on a periodic, slip or hanging-node system the prescribed DOFs are mapped into the reduced space the sweep runs in.

Cost when ω ≠ 1: one extra full residual evaluation per block per sweep.

Under bounds, over-relaxation steps past the sub-solve's answer — which is feasible by construction, while the extrapolation need not be — so the driver takes the box projector from the bounds wrapper and clips. That deviates from the paper, which backs the scalar ω off to the largest feasible value; clipping is componentwise, also feasible, and keeps more of the step.

The trade is the convergence rate, and it is not small. Alternate minimization converges linearly where Newton is quadratic, so it can need hundreds of sweeps near a propagating crack — hence the max_sweeps=200 default. It buys robustness, not speed: where Newton converges, Newton is the better choice (Farrell & Maurini, CMAME 312, 2017, compare the two directly). Sweeping is Gauss-Seidel, so the order matters, and every field block must be listed — an unlisted field's equations would never be solved, which is rejected rather than skipped. Each field is solved alone unless it is listed in a group, which is solved together inside one sweep (see "Groups" above).

Differentiable in the ordinary way: at convergence the full residual is zero, so the sweep is just a way of finding that root and lax.custom_root supplies the gradient from the full Jacobian — the alternating structure is absent from the derivative by construction.

Sparse-direct Newton — jno.solve.newton(direct=True)

An iterative linear step — the default's Jacobi-BiCGStab on the assembled tangent, or BiCGStab on the JVP with direct=False — stalls on an indefinite / ill-conditioned tangent with no good preconditioner — a Taylor–Hood velocity/pressure saddle, a stiff Carman–Kozeny phase-change drag in a melt pool. direct=True instead assembles and factorizes the tangent each step with a sparse LU (the transient stepper factorizes the backward-Euler step tangent M/dt + ∂R/∂u; the steady path ∂R/∂u). It composes wherever the assembler provides that tangent — fem.solve(nonlinear=jno.solve.newton(direct=True)) on a native nonlinear problem, steady or the transient march — and stays differentiable: implicit differentiation uses a direct, transposable tangent solve on the tangent assembled at the root (the adjoint solves Jᵀ directly too, not with a stalling Krylov). damping / line_search apply unchanged. It needs the assembled tangent, so it does not apply to the matrix-free-only paths (a coupled-residual wrapper, complex) — those fail loud.

The tangent it factorizes is always the element-scattered sparse one, in 1-D as in 2-D/3-D. 1-D used to build its nonlinear tangent with a global jax.jacfwd instead, on the assumption that only the then-default matrix-free Newton would ever ask for it; direct=True does ask, from inside the Newton loop, where a dense tangent cannot be sparsified at all (BCOO.fromdense needs a concrete nse). Besides working, the scattered tangent is O(nnz) rather than O(N²) — which is what 1-D node counts want.

A direct linear= slot selects it. lu, dense and amg all need an assembled matrix, so pairing one with the matrix-free Newton has nothing to factorize. fem.solve(linear=jno.solve.lu(backend="host")) on a nonlinear or transient problem therefore routes to the direct Newton, and that slot is the solver that runs on the assembled tangent (and on Jᵀ in the adjoint); precond= materializes against the same assembled operator. Which factorization you pick is not cosmetic here: on a 26-step Rayleigh–Bénard march (three fields, saddle, nonlinear) the matrix-free Newton that was the default when this was measured (direct=False now) takes 20.1 s, linear=jno.solve.lu() 7.6 s and linear=jno.solve.lu(backend="host") 3.1 s, all to the same 2.8e-07 per-step Newton residual. An explicit matrix-free nonlinear= alongside a direct linear= is contradictory and raises rather than picking one silently.

Limits, since they are not obvious. This routing only fires when the linear slot is direct. A precond= that needs an assembled matrix (jacobi, an unbuilt amg) gets one from the default Newton; on the matrix-free routes (newton(direct=False), picard(), a problem with no assembled tangent) it raises where it is materialized. picard() has no assembled-tangent form: its linearization is the lagged matrix-free JVP, so nonlinear=picard(), linear=lu() raises — moving to newton(direct=True) there changes the algorithm, not just the solver. And the direct Newton needs the assembler to supply a tangent, so the matrix-free-only routes (a coupled-residual wrapper) fail loud either way.

Reusing the factorization — newton(direct=True, reuse=True). Lagged-Jacobian (chord / Shamanskii) Newton: keep stepping against the last factorized tangent while each step still halves the residual, refresh it when the contraction falls below that, and reject a stale step that does not reduce the residual at all (Kelley, Solving Nonlinear Equations with Newton's Method, SIAM 2003, §1.5, §2.4). It trades expensive factorizations for cheap extra steps, so it only pays where a factorization costs many residual evaluations, and only with a backend that keeps its factorization — lu(backend="cudss"), "pardiso" and "host" cache on the matrix values, so the same tangent again is a solve; backend="device" refactorizes anyway. fem.stats["nonlinear"]["factorizations"] reports the count on an eager solve.

Measured, honestly mixed. A 2-D manufactured nonlinear Poisson: 5 → 2 factorizations, same root. A 3-D rigid-plastic rolling model (22k DOFs, cuDSS, RTX 3070): the continuation ladder dropped from 137 to 75 factorizations, but the run got slower (134 s against 120 s) — there a factorization costs ~0.12 s while recompiling for each remeshed geometry costs ~6 s, so the linear solve was never the bottleneck. Measure where your time goes before reaching for it.

Writing your own

User extension is duck-typed — a linear solver is any fn(A, b, *, M=None, x0=None) -> x with A a jno.solve.LinearOperator (.mv, .T, .diag(), .bcoo, .dense()); a preconditioner is any ctx -> (v -> M⁻¹v):

def my_precond(ctx):                      # ctx.A, ctx.diag(), ctx.fem
    inv = 1.0 / ctx.diag()
    return lambda v: inv * v
u = fem.solve(linear=jno.solve.cg(), precond=my_precond)

Calling a solver directly (outside fem.solve) takes the preconditioner as M=, which is the application v -> M⁻¹v. A jno.precond.* spec is accepted there too and materialized against A on the way in — jno.solve.cg()(A, b, M=jno.precond.jacobi()). A bare callable is always the applier, never a ctx -> applier factory: as a precond= slot it would be the factory, nothing about a callable tells the two apart, and guessing wrong would apply a preconditioner you did not ask for. Specs that need eager preparation (form, which assembles an auxiliary operator) require fem.solve(precond=...) — a direct call has no owning FEM.

solve_fn= is the level below that: instead of a slot, you take the assembled system and return the answer. jNO does nothing to it — no firewall, no convergence gate — so it is the escape hatch for a solver that fits none of the shapes above:

u = fem.solve(solve_fn=lambda A, b: jnp.linalg.solve(A.todense(), b))   # dense LAPACK
# max|u - default| = 2.9e-10 on a 142-DOF Poisson operator

A arrives exactly as assembled — a BCOO, not a densified copy. That is deliberate: densifying is O(n²) on an O(nnz) operator (20.5 GiB against 9.4 MiB at n=51,843), so it would make bringing a sparse solver pointless at the sizes where you would want one. Call .todense() yourself if that is what you need, as above.

If your callable is pure JAX it inherits jit/vmap/AD automatically. On the matrix-free nonlinear path (newton(direct=False), picard(), or no assembled tangent) the precond spec is materialized per Newton/Picard linearization against the JVP operator — so form, inner(...), chebyshev, a pre-built amg, and their block_diag/triangular compositions all work; only specs that need the assembled matrix (jacobi, an unbuilt amg) raise. The default Newton materializes the spec against the assembled tangent instead.

Diagnostics — what the solver actually did

What a solve prints

Every build and solve logs one line each for what was built, which solver ran, and what came back:

INFO: Preprocessed mesh connectivity: 85 points, 133 triangles · h 0.0927–0.141 · worst aspect 1.38
INFO: FEM(linear, dofs=85, terms=['volume', 'dirichlet@boundary'])
INFO: solve: linear · bicgstab (matrix-free) + jacobi
INFO: solved: linear · 0.115 s · u in [0, 0.06101]

• FEM(...) is the mode fixed at build and every term as jNO classified it. A boundary condition that landed nowhere is missing here.
• solve: is the route actually taken, including defaults you did not pass.
• solved: is the solution range, with ALL ZERO — is the load term present? when the right-hand side was empty. A march adds its steps and the step that came closest to its tolerance:

INFO: solved: nonlinear · 0.454 s · load-path march · 40 steps over τ ∈ [0, 1] · tightest step 32 (τ=0.7949) at 82.9% of its bound · u in [0, 0.8375]

The lines cost the solve nothing it does not already pay. Nothing is added inside a traced computation: a jitted transient march compiles to the same program with or without them. The solved: line reduces the solution on the device and moves four numbers to the host (0.3 ms against a 63 ms solve at 73k dofs on a GPU). It skips itself when the solve is traced (jit/grad/jno.core). It does not compute ‖Au − b‖: that is a sparse matvec, 10 % of the same solve, and the solve already checks its own residual and raises (below). Ask for it with jno.info(sol, context=fem). gmsh's own progress output is suppressed; its warnings are logged as WARNING: gmsh: ….

fem.stats

After any fem.solve(), fem.stats reports what happened without changing the solve's return: mode, dofs, wall_s (dispatch time — JAX is async; block on the result for compute time), the linear/precond slot reprs, nonlinear (driver, final residual norm against its bound, converged flag, and the step count — every driver reports it), linear_iterations (below), and amgx_cache occupancy when jaxamg served the solve. Populated on eager paths; a solve wrapped whole in jit/vmap/grad records the slots but no residuals — reporting needs a concrete value. The convergence check below is not so limited: it fires under a transform too.

sol = fem.solve(nonlinear=jno.solve.newton(direct=True), linear=jno.solve.lu(backend="host"))
fem.stats
# {'mode': 'nonlinear', 'dofs': 44, 'wall_s': 0.31, 'linear': 'jno.solve.lu-host(...)',
#  'precond': None, 'nonlinear': {'driver': 'newton_direct', 'residual': 1.6e-07,
#                                 'bound': 1.7e-06, 'steps': 3, 'converged': True}}

solve_index counts solves on the form (the first one includes tracing and compilation), and a solve that raised still writes fem.stats, with an error entry. It used to leave the previous solve's stats in place, and those would read as this solve's.

Krylov iterations — fem.stats["linear_iterations"]

How many iterations each linear solve took — the number that says whether a preconditioner works:

fem.solve(nonlinear=jno.solve.newton(direct=False))
fem.stats["linear_iterations"]
# {'solves': 4, 'iterations': [52, 62, 59, 54], 'total': 227, 'max': 62,
#  'by': ['the newton_krylov inner BiCGStab']}          # one count per Newton step

iterations lists each forward solve's count in the order they finished, and transpose the adjoint solves' when a gradient ran any. breakdown: True flags a BiCGStab that stopped on one. The default solve, cg, bicgstab, minres, fgmres, cocg, chebyshev, the matrix-free Newton's inner BiCGStab and the assembled-tangent Krylov all report. Upstream restarted gmres exposes no count and is listed under uncounted; a direct factorisation reports None. Counts are recorded for an eager, steady, non-adaptive solve. A march records none per step — its compiled step makes no host call, and fem.stats["march"] is its record — and neither does a deferred node evaluated later under jno.core.

cg and bicgstab are jax.scipy's own loops, keeping the iteration count they discard: the same arithmetic (x identical on CPU, within 1e-10 on GPU from reduction order) and the same time. The count reaches the host through one callback per solve, a fixed ~0.08 ms (CPU) / ~0.15 ms (GPU): +7–11 % on a 790-DOF solve, unmeasurable at 46,691 DOFs.

A march records every step — fem.stats["march"]

A march (a tau= load path, an arc-length path, continuation=, a transient) records what each step did. grid is the coordinate at each step. residual and bound are each step's final residual and the tolerance it was judged against. step_s is each step's wall time:

march	per-step residual	per-step time
load path, arc-length	✅	— one compiled lax.scan: only the mean wall_s / steps
continuation=	✅ nonlinear	✅ every rung is materialised, so its time is real compute
transient, nonlinear	✅ once evaluated	— one compiled lax.scan
transient, linear	—	— one compiled lax.scan

A step inside a lax.scan is not a host-visible event, so no per-step time is reported for it rather than an invented one. A transient solve returns a deferred node: its record says deferred=True with the step count, dt and window until the node is evaluated eagerly (.fn()), which adds the evaluation's wall_s and an evaluation count. Under jno.core the march runs inside a compiled program and nothing is recorded.

A march that fails keeps its record up to and including the failing step, at the same 1-based index the error names. jno.info(fem) renders it:

  march
    kind         load-path march
    steps        12 over τ ∈ [0, 1]
    time         542 ms for the solve · mean 45.1 ms/step — includes compilation
                 one compiled lax.scan: a step has no wall time of its own, only the mean
    convergence  all 12 steps converged · tightest step 11 (τ=0.9091): residual at 45.6% of its bound

"Tightest" is the step that came closest to its bound, which is the first place to look when a finer grid or a harder load is about to break the march. jno.info(fem, deep=True) lists every step.

A solve that did not converge raises — including the adjoint

Every linear solve is residual-checked against its own operator before it returns, and so is its transpose — the plain fem.solve() default included. That matters because a Krylov iteration leaving on its step cap returns the last iterate with no signal at all, so a broken solve is otherwise indistinguishable from a good one:

jax.grad(loss)(theta)
# RuntimeError: jno.solve.fgmres did not solve the system: relative residual 2.8e-02 against a
# 1e-04 gate. This is the ADJOINT (transpose) solve, not the forward one -- the solution itself is
# fine and only the GRADIENT is affected. A Krylov method can break down on A^T where it converges
# on A, so reach for a transposable-robust inner solver (jno.solve.gmres(), jno.solve.lu()) or a
# preconditioner rather than loosening the tolerance.

Where the refusal surfaces — dispatch is asynchronous

The gate is evaluated inside the traced solve, and reports through a host callback. JAX dispatch is asynchronous, so that callback has typically not run when the solve returns: the refusal appears only once something materialises the result. Measured on the same failing solve, the raise reaches the caller on the CPU backend (callbacks run inline) and is swallowed into a log line on GPU.

So the failure is recorded first and raised at a boundary that is guaranteed to exist:

you called	when it raises
fem.solve(...)	at the call — the result is materialised and the gate drained before returning
an eager jno.solve.cg()(A, b)	at the call, same way
your own jax.jit / jax.grad around a solve	jNO has no post-execution hook there. The failure is recorded, and the raise from the callback is best-effort — it lands on some backends and not others

The last row is a real limit, not an oversight: once the solve is inside your transform, the library gets no callback after it executes. If you need certainty there, materialise inside the transform-free part and let fem.solve do the checking.

Why the adjoint gets its own check

The adjoint is the half nothing used to check. lax.custom_linear_solve's transpose_solve had no convergence test, and the eager test on the forward solve is a no-op under tracers — so the guard was off precisely when gradients were being taken.

Measured on a 2-D strong-form Poisson operator: BiCGStab converged on A (relative residual 2.9e-05) and diverged to 5.5e+26 on Aᵀ for the same problem. The last iterate came back as a gradient wrong by twenty orders of magnitude, and nothing anywhere raised.

Cost is one extra matvec per solve. The host round-trip that reports the failure sits in the failing branch, so a converged solve never pays for it: 0.32 ms against 0.24 ms per call on a small dense GMRES, and no measurable difference on the reverse pass. The fem.solve() default got faster — 1.90 ms against 2.73 ms at 1441 DOFs — because checking from inside the trace replaced an eager check that cost its own matvec plus a host sync.

A transient march judges every step, not just the last one

A nonlinear transient step is a Newton solve, and every one of them runs inside the march's lax.scan — precisely where the driver's own convergence check disables itself, because it needs a concrete residual. So the step carries its residual norms out of the scan and fem.solve() tests them where they are concrete, against the driver's own rtol/atol:

fem.solve(nonlinear=jno.solve.newton(direct=True, rtol=1e-6, atol=1e-6))
# RuntimeError: fem.solve: the transient march did not converge at step 36 of 120 (t=0.00036):
# residual norm 3.658e+02 against the tolerance atol + rtol*||r(u_prev)|| = 1.011e-02
# (atol=1e-06, rtol=1e-06). That step is NOT a root, and every later step inherited it as its
# starting state — the whole trajectory past this point is unreliable. Globalize the per-step solve
# (jno.solve.newton(line_search=True), or damping<1), take smaller time steps (a finer
# domain(time=(...)) grid), or raise the driver's max_steps.

This is the same check the load-path march already applied, on the same numbers; only the label and the advice differ. It covers the default θ march, jno.solve.bdf2() and jno.solve.sdirk() (judged per stage: every stage is a Newton solve).

Why a silent transient failure is worse than it sounds

Two properties make a capped march genuinely hard to spot, both measured on a coupled melt-pool model (laser + phase change + Marangoni convection):

• It is not reproducible. Two runs of the same command reported peak melt velocities of 2.98 and 33.5 m/s — an unconverged iterate is whatever the step cap happened to leave.
• A diverged march is faster than a healthy one. Once the residual is NaN the loop condition ||r|| > tol is False, so Newton exits on its first iteration and every remaining step is nearly free. The usual "it got slow, something is wrong" signal is inverted: the run that finished quickly was the broken one, and it returned 2e+200 K with nothing raised.

Cost is two residual evaluations per step, and only on a nonlinear block — a linear step is a linear solve with its own gate and is not judged here.

Under a transform the guard is off

Inside jax.jit/grad/vmap of a whole march the norms are themselves traced, so the test cannot concretise and no-ops — the same trade every eager check in jNO makes. There, the driver's max_steps is all there is.

A preconditioner is exempt — it is inexact on purpose

jno.precond.inner(jno.solve.cg(tol=1e-2, maxiter=30)) asks for two digits deliberately; that is the whole reason the outer solver must be flexible (jno.solve.fgmres()). The gate is a property of the call, not of the solver spec: the same cg() is checked as an outer solve and unchecked as a preconditioner application. The convergence contract belongs to the outer solve, and that one is still gated.

Where the time went — fem.solve(profile=True)

fem.stats says what the solver did; profile=True says how long and where. It runs the solve inside a JAX Perfetto trace, prints the DOF count and wall time, and writes the trace to ./jno_traces for chrome://tracing or ui.perfetto.dev:

u = fem.solve(profile=True)
# [fem profile · 75 DOFs · linear]  wall 82.8 ms  ·  Perfetto trace → ./jno_traces

The same flag exists on jno.core(...).solve(profile=True) (see Evaluation) and on jno.rcwa(...).solve(profile=True).

Profile a concrete solve

It runs eagerly, at the current parameter values. A parametric solve returns a deferred trace node with no numeric work to time, so there is nothing to profile — and ./jno_traces is not in the repo's .gitignore, so the trace directory will show up as untracked.

Transient problems

What a march keeps — only the frames you ask for

Evaluated eagerly (fem.solve(...).fn(), fem.solve(k=2.0)), a march keeps no past on the device. Each step writes its state into a save slot when a frame at save_ts will read it (a save time on the grid, or the two grid points either side of one off it) and into a scratch row otherwise, and the frames come back as a host NumPy array. Saving every step of a march whose frames would outgrow 1/16 of the device runs it in chunks: each chunk's frames are copied out while the next chunk is already queued, with nothing inside the compiled step (no callback, no sync). The answer is bit-identical to the one-scan march on CPU, under every scheme below; on a GPU it moves only by the run-to-run noise the old march had (measured 4–7e-11, scatter-add order).

Measured (RTX 3070, 2-D linear heat, 22,801 DOFs, 5 frames saved): peak device memory 186 → 59 MiB at 500 steps and 708 → 59 MiB at 2,000 steps; wall time unchanged.

Pass save_ts= for a long march. The default saves every grid point, which keeps the whole trajectory — on the host now, not the device. Under jit/grad (jno.core, an inverse problem) the march stays one lax.scan holding every step: the adjoint needs the states.

What a march's Newton keeps — its tangent

A nonlinear march's default per-step Newton is jno.solve.newton(reuse=True): it keeps the assembled step tangent while that still contracts the residual (‖r_new‖ < ‖r_old‖/2), refreshes it the moment it does not, and rejects a step that fails to reduce the residual -- so a stale tangent costs at most one wasted solve, never a divergence. Evaluated eagerly, the march also carries the tangent from step to step (on the step-merge plan's fixed sparsity pattern), because a march's tangent changes little between steps. The step's convergence record reuses the norms Newton already computed instead of evaluating the residual twice more. And on a kept tangent the inner Krylov solve is inexact: a step contracts the residual by the tangent's own contraction, not by the accuracy of the linear solve, so each correction is solved only to a forcing tolerance -- half the outer target over the current residual, floored at 1e-3, capped at 0.1 (Dembo, Eisenstat & Steihaug 1982) -- instead of to 1e-10.

Measured (RTX 3070, the Taylor–Green LES of the turbulence tutorial: residual-based VMS + Vreman, P1/P1, BDF2): at 55,296 DOFs one tangent for an 8-step march instead of two per step, 0.80 → 0.30 s per step; at 442,368 DOFs, where the Krylov solves dominate once the tangent is kept, the inexact solves take 1,200 → 160 matvecs per step, 2.14 → 1.50 s per step -- same trajectory (energy identical to 6 digits). The answer moves only within the Newton tolerance: a reused tangent converges linearly, so the final residual lands under the test rather than far below it. Pass nonlinear=jno.solve.newton() for a fresh tangent at every iteration. Under jit/grad the march keeps the tangent within a step but not across steps (the reverse pass would store one tangent per step).

Time schemes — fem.solve(time=…)

Scheme	Order	Stability	Use
jno.solve.theta(1.0)	1	L-stable	the default — robust, damps everything
jno.solve.theta(0.5)	2	A-stable only	smooth problems; rings on stiff modes
jno.solve.bdf2()	2	L-stable	stiff / saddle-point systems — flow's workhorse
jno.solve.sdirk(order=2\|3)	2 or 3	L-stable, stiffly accurate	higher order where BDF2 is not enough; state-dependent mass allowed
jno.solve.rosenbrock("ros34pw2"\|"ros2")	3 or 2	L-stable	nonlinear blocks without Newton: one linear solve per stage
jno.solve.exponential()	exact in time	unconditional	linear, autonomous blocks
…​.adaptive(rtol=…)	wraps a one-step scheme	—	multi-rate transients

bdf2 exists because theta cannot be second order and L-stable at once. Crank–Nicolson's amplification factor tends to -1 for a stiff mode, so the mode does not decay — it alternates in sign. Measured on a heat problem whose initial condition is incompatible with its boundary (8 steps, T = 0.5, so the exact field is long dead): Crank–Nicolson reaches min u = -0.70 and still carries 0.33 at the end; BDF2 reaches -0.023.

Rows with no time derivative (Dirichlet rows, a pressure, an FDM flux row) are constraints, and every θ-scheme imposes them at the new time, as backward Euler does. The θ-average would make Crank–Nicolson satisfy them only on average over the step (a boundary started off its value would flip sign every step and never settle), and would leave forward Euler with a singular step. On a Navier–Stokes saddle system, where the pressure has no time derivative at all, that ringing is exactly what you do not want.

The first BDF2 step is plain backward Euler — a multistep method has no second level to start from. A state-dependent mass c(u)·u_t is marched in BDF2's non-conservative form c(uⁿ⁺¹)(3uⁿ⁺¹ − 4uⁿ + uⁿ⁻¹)/(2Δt), second order in time (measured on a manufactured c(u) = 1 + u²); it is not the conservative form an enthalpy mass H(u)_t would want. It composes with periodic ties (the march carries the reduced state; the mass action is reduced like the residual). Refused loudly rather than mis-integrated: a second-order-in-time (u_tt) block (assembled at θ=½ so that it is not damped), and .adaptive() (step doubling sizes a one-step method).

Every scheme needs each transient term linear in u_t — c(u)·u_t, or u_t inside a contraction such as τ (u·∇v)·u_t (the time derivative inside a residual-based stabilisation). A term quadratic in u_t, or u_t inside a nonlinear function, is refused: it used to be marched with (u − u_prev)²/Δt where /Δt² belonged. The u_t⊗u_t piece of a residual-based VMS Reynolds stress is such a term — drop that piece.

Higher order — sdirk and rosenbrock

jno.solve.sdirk(order=3) is Alexander's 3-stage, order-3 SDIRK method (order 2: 2 stages; SIAM J. Numer. Anal. 14(6), 1977). Each stage is one backward-Euler step of size γ·dt from a shifted start state, so everything the default step does carries over: Dirichlet and zero-mass rows, solver slots, Newton on a nonlinear block, the complex path. All stages share one step operator, so a constant one is built and preconditioned once. A state-dependent mass c(u)·u_t is integrated at full order in its non-conservative form.

jno.solve.rosenbrock() is linearly implicit: every stage is one linear solve with the same matrix W = M + γh·J, J the tangent at the step's start — no Newton loop on a nonlinear block. The default "ros34pw2" (Rang & Angermann, BIT 45, 2005) is order 3, stiffly accurate and consistent for index-1 DAEs; "ros2" (Verwer et al., SIAM J. Sci. Comput. 20(4), 1999) is order 2. A time-varying Dirichlet row lands exactly on g(tₙ₊₁) under the default, which is stiffly accurate; "ros2" is not, and holds it only to its order (g = e^{-2t}: 5.3e-4 off at T = 0.5 after 16 steps, 1.4e-3 at the first step — see the measured orders). Re-imposing g after each step would make the wall exact, but was measured to leave an O(dt) error in the interior, so it is not done. Rosenbrock refuses a time-, parameter- or state-dependent mass (use sdirk or bdf2).

sol = fem.solve(time=jno.solve.sdirk(order=3))
sol = fem.solve(time=jno.solve.rosenbrock())                 # ros34pw2
sol = fem.solve(time=jno.solve.rosenbrock().adaptive(rtol=1e-4))

Both expose a stepper, so .adaptive() works, and both refuse second-order-in-time (u_tt) blocks. On a nonlinear heat problem, ROS34PW2 reached a given accuracy in the least wall time of the schemes compared (θ, BDF2, SDIRK, ROS2). Time-to-accuracy depends on the problem, so compare on yours.

The slots configure the per-step solves of the default theta-method integrator: linear/precond see the step operator M + θ·dt·A — when it is time-independent the step matrix is formed once and the preconditioner materialized once before the time loop — and nonlinear drives each implicit step of a nonlinear block. Second-order-in-time (u_tt) flows through the same augmented block. Each step warm-starts from the previous state (so x0= is rejected).

A constant step operator factorizes once, not once per step

The step matrix is formed once (above), and the host factorization is cached on the operator's content, so a march that solves against the same matrix every step pays one factorization for the whole trajectory — and the transpose solve reuses it, so the adjoint pass adds none. On a 51-step heat march that is worth 1.55× of whole-solve wall clock at 8,355 DOFs and 2.9× at 23,934. The gain grows with mesh size, because factorization cost grows faster than the per-step solve; at 513 DOFs it is 1.05×, where the factorization is sub-millisecond.

What it does not help: a nonlinear march, or any Newton loop. There the tangent's values change every iteration, so every call legitimately misses and pays a content hash (~1–2% of a factorization) for nothing. Reusing only the symbolic analysis — which genuinely is constant when the sparsity pattern is fixed — is not something SuperLU exposes through scipy. It is also host-path-only: lu() on GPU goes through cuSolver's spsolve, a single fused call with no factorization object to keep.

Adjoint memory — the march is gradient-checkpointed

Reverse-mode through the default integrator rematerializes each step in the backward pass instead of storing every step's solver internals for the whole trajectory. Measured at 8,355 DOFs × 399 steps: peak memory 968 → 112 MB (8.6×) for a gradient cost of +60% (2.98 → 4.76 s), gradient identical to 10 digits. The trade is deliberate: a differentiable march OOMs long before it is time-walled on consumer cards. A pure forward solve is unaffected — checkpointing is the identity outside differentiation.

Complex problems

Complex problems are assembled as one real 2n system over the stacked [Re; Im] state — the real-equivalent block [[A_r, -A_i], [A_i, A_r]] — at assembly rather than at solve time. A complex transient is therefore an ordinary transient block (the slots configure its per-step solve as above, and theta / adaptive / exponential all apply), and a complex steady problem is an ordinary linear system: the linear / precond slots and x0= work on it, with a complex x0 mapped into the block layout for you. Its default solver stays sparse-direct, not the matrix-free BiCGStab real elliptic systems get — the real-equivalent block is indefinite for Helmholtz/PML, where Jacobi-BiCGStab does not converge.

One exception keeps a dedicated path: a complex-native preconditioner (ams) solves A_r + i·A_i directly rather than the block, so the Re/Im legs are retained for it.

Bloch (quasi-periodic) ties

A Bloch (quasi-periodic) tie fuses like everything else. Its complex prolongation P cannot reduce the Re/Im legs independently, but on the fused [Re; Im] state the same tie is the real prolongation B(P) = [[P_r, -P_i], [P_i, P_r]], and the ordinary real congruence B(P)ᵀ A B(P) equals the Hermitian reduction P^H A_c P the Bloch space requires. Consequences: solve_fn=, the linear/precond slots and x0= all apply to a Bloch problem (each used to be silently discarded by a dedicated block routine); a Bloch tie composes with a complex transient (the quasi-periodic plane-wave march, previously a dtype crash); and a real weak form with a Bloch tie is promoted to the complex path automatically — the phase makes the field complex anyway, and the real path's bilinear Pᵀ A P is not a Galerkin projection for a complex P (measured 8.1 rel-L2 off the Hermitian answer on a manufactured mode, with the tie itself satisfied exactly).

A coupled (multi-field) complex steady system — coupled Helmholtz-type equations — takes the same Re/Im split through the coupled assembler: one fused real 2n block over [Re_all; Im_all], with fem.offsets still listing the per-field blocks of the recombined complex solution. Scope: steady and linear (a complex nonlinear coupled form and a complex coupled transient refuse, as everywhere).

Essential values on a complex form must be real. The two legs share one Dirichlet row set, which imposes Re u = g with Im u = 0 — right for a real g, and the usual case (the complexity lives in the operator and the source). A complex g is not expressible there: pinning Im u = g_i would need the imaginary leg's rows zeroed rather than set to identity, and the symmetric elimination's known-column lift is cross-leg (the real equation needs A_r[:,j] g_r - A_i[:,j] g_i, which no per-leg elimination produces). It raises a clear error. Carry the complex part in the operator or the source.

adapt= composes with a complex transient too: the stacked [Re; Im] halves transfer across each remesh as a doubled field layout, the modulus |u| drives the remesh metric (refining on Re alone would miss a rotating phase), and the saved frames come back complex. A criterion= is refused there by name: it would be assembled against the stacked real state, whose halves it cannot tell apart. Drop it to remesh on the modulus.

Not yet supported (clear errors): a Bloch tie on a real transient march (the phase forces a complex field — make the problem complex, or use a plain tie) or on a nonlinear form (complex nonlinear forms are not wired).