# nlls_gram

> Metric-aware Levenberg-Marquardt nonlinear least squares for JAX pytrees, aimed at interpolation problems (zero-residual roots where minimum-metric-norm selection matters, wide or tall). 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). 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 (assembled pipelines) and metric_solve_dtype (metric callbacks).

## Docs

- [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 — tridiagonal/state-space-Matérn/shifted-matvec/diagonal/block-diagonal metric constructors; 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/)
