# nlls_gram

> Levenberg-Marquardt nonlinear least squares for JAX pytrees, aimed at interpolation problems (zero-residual roots where minimum-norm selection matters, wide or tall). Two solvers share one init/update/solve protocol: RidgeLevenbergMarquardt minimizes ||r||^2 + ridge*||x_m||_W^2 for a positive-definite Metric W on the metric block x_m of x = [x_m; x_f] (free block x_f unpenalized, n_f = len(x) - metric.size inferred at init) with the ridge weight as annealable traced state (ridge_continuation), putting the minimum-seminorm/min-RKHS-norm selection in the objective (nonlinear Tikhonov). The metric is supplied through factor callbacks for W = F'F (factor_apply/factor_solve/factor_solve_transpose/norm, each receiving a MetricContext with the live solver state; IdentityMetric is plain ridge, RepeatedFactorMetric the kernel workhorse taking F = jnp.linalg.cholesky(K, upper=True)) and the solver runs entirely in the whitened variable y = F_bar x with constant penalty rows [I | 0] — spectral floor at the ridge, Cholesky()-path accuracy at deep ridge, lambda-free factorization so continuation composes. Typed solver configs Cholesky()/QR()/CG(preconditioner, ...) with ad_solver Cholesky()/CG(...); the CG preconditioner is a typed Preconditioner subclass (apply(v, damping, ctx); IdentityPreconditioner() opts out; BlockEigenPreconditioner applies damping-analytic block-eigenbasis state carried in the residual args and rebuilt adaptively from solve callbacks via block_eigen_state), required in both roles. GN-implicit AD posed on y; conjunctive gtol+atol stopping in the whitened geometry (steps in the W-norm, gradients in the dual W^{-1}-norm; info.penalty_grad_norm = sqrt(penalty_value), so gtol ~ 1e-3*ridge*sqrt(q)); runs at the residual dtype, QR() the extreme-conditioning fix. LevenbergMarquardt is metric-damped LM where selection comes from the damping geometry (its GramMetric is a separate contract). Per-step update(), a jitted solve() loop with resettable hyperparameters and callbacks, implicit differentiation with respect to external parameters p, and pluggable parameter-space metrics (kernel/RKHS). Metric-LM linear solvers: a shape-adaptive dense default (linear_solver="auto" resolves to gram_cholesky when n > m, else normal_cholesky), qr/augmented_qr, and matrix-free gram_cg/normal_cg/lsmr; jacobian_mode="auto"/"fwd"/"rev" assembles the dense Jacobian from the small side. Implicit differentiation has explicit methods: direct, svd, qr, augmented_qr, gram_cg, normal_cg, and regularized_normal_cg. ad_solver="auto" selects direct for square systems, the matching CG space for nonsquare CG systems, and svd otherwise. Targeted float64 promotion via linear_solve_dtype (metric-LM assembled pipelines) and metric_solve_dtype (metric callbacks).

## Docs

- [Ridge LM — the ridge objective on the metric/free block split, selection theorem, the whitened change of variables y = F_bar x (W = F'F, augmented stack [J~; sqrt(ridge)[I 0]]), continuation schedule, the two-phase stopping picture and the whitened gtol calibration recipe (gtol ~ 1e-3*ridge*sqrt(q); noise-floor regime), the Metric interface (factor ops + MetricContext, exact-factor and identity-hashing contracts, free-block guidance), typed Cholesky()/QR()/CG(...) solver table with typed Preconditioners, GN-implicit AD contract, migration from the metric formulation](https://highdimensionaleconlab.github.io/nlls_gram/ridge_lm/)
- [Tuning guide — read first for solver selection and hyperparameter heuristics](https://highdimensionaleconlab.github.io/nlls_gram/tuning_guide/)
- [Main docs — API contracts, math, linear solvers, performance notes](https://highdimensionaleconlab.github.io/nlls_gram/)
- [Callbacks and cookbook — solve-loop callback contract, resettable hyperparameters, recipes](https://highdimensionaleconlab.github.io/nlls_gram/callbacks/)
- [Multi-start — retry failed solves from fresh draws or race starts in parallel under vmap; draw/accept hooks, key schedule, winner-only implicit AD](https://highdimensionaleconlab.github.io/nlls_gram/multi_start/)
- [Metrics — Metric callback contract, cholesky helper, dense and matrix-free examples, and the iterate-aware MetricFactory (prepare/build) that rebuilds the metric from the current iterate and residual aux each accepted step](https://highdimensionaleconlab.github.io/nlls_gram/metrics/)
- [Utilities — dense, diagonal, repeated-shifted dense, and repeated-shifted state-space metric constructors; the Matérn state-space producer; Sherman-Morrison, Woodbury, identity, and randomized Nyström (FTU) preconditioners serving the CG hooks; the normal-space preconditioner contract for linear_solver="normal_cg" (SPD approximation of (B'B + damping I)^-1 that must preserve range(B') on rank-deficient problems); the iterate-adaptive PreconditionerFactory (gram_cg-only) that rebuilds the dual preconditioner from the current iterate each step; the dual-only padded-residual helper; matrix-free LSMR (linear_solver="lsmr") for the whitened damped subproblem when the squared Gram/normal operators are ill-conditioned at small damping, with an optional WhitenedPreconditioner parameter-space right-preconditioner that clusters the LSMR spectrum without changing the damped subproblem; and Krylov recycling (RecycleConfig / deflated_pcg, gram_cg-only) that carries a harvested deflation basis across LM steps](https://highdimensionaleconlab.github.io/nlls_gram/utilities/)
- [Implicit differentiation — JVP and automatically transposed VJP with respect to p; MAX_STEPS is usable by default with an opt-in strict policy; zero failed-solve derivatives from a JVP-safe initial point; shape-first auto dispatch; direct square solves; SVD pseudoinverse; loud unregularized QR; explicit augmented-QR and regularized-normal-CG ridges; matrix-free Gram/normal CG; no default regularization; aux outputs](https://highdimensionaleconlab.github.io/nlls_gram/implicit_ad/)
- [Metric Gauss-Newton — minimum-norm math and kernel metric choices](https://highdimensionaleconlab.github.io/nlls_gram/gauss_newton/)
