API

Solvers

nlls_gram.RidgeLevenbergMarquardt

Bases: LevenbergMarquardtBase

Levenberg-Marquardt for the ridge objective F(x) = ||r(x, args, p)||^2 + ridge * q(x) with q(x) = ||x_m||_W^2 over a JAX pytree x, built for underdetermined interpolation problems where the minimum-seminorm root argmin q s.t. r = 0 is the target.

The flattened parameter vector splits as x = [x_m; x_f]: the METRIC BLOCK x_m (the leading metric.size coordinates, penalized by the positive-definite :class:~nlls_gram.Metric W) and the FREE BLOCK x_f (the remaining n_f = len(x) - metric.size >= 0 coordinates, unpenalized -- the full-space penalty is PSD by design). The metric is supplied through factor callbacks for W = F'F (:class:~nlls_gram.IdentityMetric is plain ridge, :class:~nlls_gram.RepeatedFactorMetric the kernel workhorse), and the solver runs entirely in the whitened variable y = F_bar x with F_bar = blockdiag(F, I_{n_f}) -- an exact, constant linear change of variables, never a matrix: only the factor ops are ever applied.

The selection lives in the objective: under free-block identification at the solution, each fixed-ridge problem has an isolated minimizer x_ridge with x_ridge -> x* = argmin q s.t. r = 0 at O(ridge) error as ridge -> 0 (nonlinear Tikhonov theory: Engl-Kunisch-Neubauer 1989; Engl-Hanke-Neubauer 1996 Ch. 10; the seminorm formulation is Elden 1982). Plain Gauss-Newton on the underdetermined system converges to some root without selecting the minimal-norm one (Campbell-Kunkel-Bobinyec 2012; Pes-Rodriguez 2022), which is why the package's alternative is metric-damped :class:~nlls_gram.LevenbergMarquardt; this solver instead makes every inner problem a well-posed NLLS. Annealing ridge per stationarity event (:class:AnnealRidge) is the iteratively regularized Gauss-Newton method (Bakushinskii 1992; Kaltenbacher-Neubauer-Scherzer 2008). For kernel metrics each inner step is a kernel ridge regression of the relinearized equations (Chen-Hosseini-Owhadi-Stuart 2021).

Everywhere in the implementation the equivalent augmented-residual view is used, posed in y: standard EUCLIDEAN LM on R(y) = [r(x); sqrt(ridge) y_m] with J~ = J F_bar^{-1}, A = [J~; sqrt(ridge) [I 0]], and step (J~'J~ + ridge E + mu I) delta_y = -(J~'r + ridge [y_m; 0]) where E = blockdiag(I_{n_m}, 0); the x-space step is delta_x = F_bar^{-1} delta_y and the iterate is stored in x. The penalty rows are affine (constant, even) in y, so geodesic acceleration (Transtrum-Sethna 2012), accept/reject, and the implicit-AD rule reduce to the standard formulas with second-derivative contributions from the penalty identically zero; the trial penalty uses the linearity F_bar(x + delta_x) = y + delta_y, so no second factor application is ever needed. damping (mu) is a plain trust-region parameter in the whitened geometry (mu ||delta_y||^2 = mu ||delta_x_m||_W^2 plus the Euclidean free block), fully decoupled from ridge (lambda): mu moves every step by accept/reject, lambda only when a callback replaces lm_state.ridge. The whitened normal matrix J~'J~ + ridge E carries a clean spectral floor at ridge on the metric block regardless of W's conditioning -- the reason the default Cholesky() path stays accurate at deep ridge.

ridge is STRICTLY POSITIVE by contract -- the constructor validates it, callbacks must keep it positive, and continuation floors are positive -- so J~'J~ + ridge E is positive definite under the identification condition at every reachable state. ridge=None resolves at init to sqrt(finfo(dtype).eps) of the residual dtype (float64 ~ 1.5e-8, float32 ~ 3.5e-4). The metric's factorization never involves the ridge weight, so continuation composes unchanged.

linear_solver picks the algebra for (J~'J~ + ridge E + damping I) delta_y = -g through a typed config (:class:~nlls_gram.Cholesky / :class:~nlls_gram.QR / :class:~nlls_gram.CG -- each method's knobs live on its own config); every path shares its factorization or inner-solve setup between the velocity and geodesic-acceleration solves:

• Cholesky() (the default): dense normal equations. J' is materialized per jacobian_mode, J~' follows by one batched factor_solve_transpose, and the assembled G = J~'J~ + ridge E (ridge added on the metric-block diagonal only) is cached across rejected steps (only the damping shift re-factors).
• QR(): MINPACK-structured damping-row QR (More 1978), stable at small ridge/damping where forming G squares the condition number. One QR of the augmented stack [J~, sqrt(ridge) [I 0] | b~] with b~ = [r; sqrt(ridge) y_m] is cached per (x, ridge) (the extra column carries the Q-transformed residual); each update re-factors [R; sqrt(damping) I] with its Q retained and solves the velocity as a backward-stable least-squares problem at cond(A), never cond(A)^2. The geodesic-acceleration RHS reuses the damped factor through corrected semi-normal equations with one fixed iterative-refinement pass (Bjorck 1987; Bjorck 1996 Sec. 6.6.5) -- the second-order correction tolerates the squared conditioning.
• CG(preconditioner, ...): matrix-free preconditioned CG on the damped normal operator J~'J~ + ridge E + damping I itself -- the same SPD system the Cholesky() path factors, with products through jvp/vjp and the metric's factor callbacks instead of an assembled G. The config requires a typed :class:~nlls_gram.Preconditioner (:class:~nlls_gram.IdentityPreconditioner opts out): its apply(v, damping, ctx) -- an SPD approximation of the damped inverse, handed the same :class:~nlls_gram.SolverContext as the metric ops -- sits in CG's M slot with the live damping. The operator carries the ridge spectral floor on the metric block, so the preconditioner only has to capture J~'J~'s structure.

Stopping (solve): a ridge solve has two phases -- the residual drops to its floor fast, then the iterate slides along the interpolation set resolving the null-space (selection) component while ||r|| stays essentially constant. A pure-residual test is blind to phase 2, so gtol bounds the whitened ridge stationarity info.grad_norm = ||J~'r + ridge [y_m; 0]|| (the dual W^{-1}-norm of the half-gradient) and xtol the accepted whitened step norm ||delta_y|| (the W-norm of the step) -- they mean "done with the current fixed-ridge problem" -- while atol is a CONJUNCTIVE filter on the TRUE residual: convergence requires (gtol or xtol fired) AND (atol == 0 or sqrt(resid_loss) <= atol). atol alone never stops the solve -- a pure-residual test would stop at step 0 from any interpolating start before the seminorm is minimized -- so solve rejects atol > 0 without a positive gtol or xtol. Calibrating gtol is clean in the whitened geometry: info.penalty_grad_norm = sqrt(q(x)), so gtol ~ 1e-3 * ridge * sqrt(q(x*)) resolves the selection to ~1e-3 relative accuracy, with sqrt(q(x*)) (the solution's seminorm) usually known to an order of magnitude before any pilot run. In a continuation run the callback keeps lowering ridge, so intermediate stationarity with a large residual correctly does not stop the loop.

solve(...).x has a custom implicit AD rule with respect to p: Gauss-Newton implicit differentiation of the ridge stationarity, posed on the whitened variable -- (J~'J~ + ridge E) y_dot = -J~'(dr/dp) p_dot then x_dot = F_bar^{-1} y_dot -- with ridge frozen (stop-gradient) at the returned state's value; the continuation schedule's and multi-start selection's dependence on p is deliberately ignored, and no damping enters the AD matrix. Exact differentiation carries two extra terms (sum_i r_i * d2r_i/dx2 in the matrix and (dJ'/dp) r on the right); both are exactly zero for an affine residual and first order in ||r|| otherwise, so the rule is exact when the converged TRUE residual vanishes. The caveat to state plainly: the first-order-in-||r|| absolute error translates to a small RELATIVE tangent error only under conditioning assumptions (J'J bounded below). On a genuinely underdetermined CURVED system the null-space block of the exact tangent carries the constraint-curvature term sum_i nu_i d2r_i with multiplier nu = r/ridge -- which does not vanish relative to the ridge-scaled penalty as ridge -> 0 -- so tangent components read off null-space directions retain an O(1)-in-ridge Gauss-Newton bias proportional to the residual curvature. ad_solver=Cholesky() assembles and factors; :class:~nlls_gram.CG runs matrix-free CG on the same operator with its preconditioner hook (PD does not mean well-conditioned -- unpreconditioned CG degrades as ridge shrinks). ad_solver=None (the default) matches the forward path: Cholesky() for the dense forwards, and a CG forward keeps its family AND its preconditioner -- the typed apply is damping-analytic, so the forward hook serves the undamped system at zero damping exactly (hooks marked requires_positive_damping fall back to unpreconditioned), while the AD tolerance and iteration budget stay at the AD defaults; pass ad_solver=CG(...) to pin those. Failed statuses return exact zero tangents for result.x/result.aux and evaluate the masked tangent program at stop-gradient copies of the caller's original inputs and the INITIAL ridge (never a possibly-invalid callback-produced value).

The init/update/solve protocol, callback contract (:class:~nlls_gram.LMContext -> :class:~nlls_gram.LMAction), multi_start, save_steps, and :class:~nlls_gram.LMSolveResult are shared with :class:~nlls_gram.LevenbergMarquardt; code written against that solver ports by changing the constructor (its damping metric -> this metric) and reading info.resid_loss where it means equation error, since info.loss here includes the penalty. Multi-start ranking uses the ridge objective at each lane's own final ridge -- comparable across lanes when they share a continuation schedule.

Source code in src/nlls_gram/ridge_lm.py

class RidgeLevenbergMarquardt(LevenbergMarquardtBase):
    """Levenberg-Marquardt for the ridge objective
    ``F(x) = ||r(x, args, p)||^2 + ridge * q(x)`` with
    ``q(x) = ||x_m||_W^2`` over a JAX pytree ``x``, built for underdetermined
    interpolation problems where the minimum-seminorm root
    ``argmin q s.t. r = 0`` is the target.

    The flattened parameter vector splits as ``x = [x_m; x_f]``: the METRIC
    BLOCK ``x_m`` (the leading ``metric.size`` coordinates, penalized by the
    positive-definite :class:`~nlls_gram.Metric` ``W``) and the FREE BLOCK
    ``x_f`` (the remaining ``n_f = len(x) - metric.size >= 0`` coordinates,
    unpenalized -- the full-space penalty is PSD by design). The metric is
    supplied through factor callbacks for ``W = F'F``
    (:class:`~nlls_gram.IdentityMetric` is plain ridge,
    :class:`~nlls_gram.RepeatedFactorMetric` the kernel workhorse), and the
    solver runs entirely in the whitened variable ``y = F_bar x`` with
    ``F_bar = blockdiag(F, I_{n_f})`` -- an exact, constant linear change of
    variables, never a matrix: only the factor ops are ever applied.

    The selection lives in the objective: under free-block identification at
    the solution, each fixed-ridge problem has an isolated minimizer
    ``x_ridge`` with ``x_ridge -> x* = argmin q s.t. r = 0`` at ``O(ridge)``
    error as ``ridge -> 0`` (nonlinear Tikhonov theory:
    Engl-Kunisch-Neubauer 1989; Engl-Hanke-Neubauer 1996 Ch. 10; the seminorm
    formulation is Elden 1982). Plain Gauss-Newton on the underdetermined
    system converges to *some* root without selecting the minimal-norm one
    (Campbell-Kunkel-Bobinyec 2012; Pes-Rodriguez 2022), which is why the
    package's alternative is metric-damped
    :class:`~nlls_gram.LevenbergMarquardt`; this solver instead makes every
    inner problem a well-posed NLLS. Annealing ``ridge`` per stationarity
    event (:class:`AnnealRidge`) is the iteratively regularized
    Gauss-Newton method (Bakushinskii 1992; Kaltenbacher-Neubauer-Scherzer
    2008). For kernel metrics each inner step is a kernel ridge regression
    of the relinearized equations (Chen-Hosseini-Owhadi-Stuart 2021).

    Everywhere in the implementation the equivalent augmented-residual view
    is used, posed in ``y``: standard EUCLIDEAN LM on
    ``R(y) = [r(x); sqrt(ridge) y_m]`` with ``J~ = J F_bar^{-1}``,
    ``A = [J~; sqrt(ridge) [I 0]]``, and step
    ``(J~'J~ + ridge E + mu I) delta_y = -(J~'r + ridge [y_m; 0])`` where
    ``E = blockdiag(I_{n_m}, 0)``; the x-space step is
    ``delta_x = F_bar^{-1} delta_y`` and the iterate is stored in ``x``. The
    penalty rows are affine (constant, even) in ``y``, so geodesic
    acceleration (Transtrum-Sethna 2012), accept/reject, and the implicit-AD
    rule reduce to the standard formulas with second-derivative contributions
    from the penalty identically zero; the trial penalty uses the linearity
    ``F_bar(x + delta_x) = y + delta_y``, so no second factor application is
    ever needed. ``damping`` (``mu``) is a plain trust-region parameter in
    the whitened geometry (``mu ||delta_y||^2 = mu ||delta_x_m||_W^2`` plus
    the Euclidean free block), fully decoupled from ``ridge`` (``lambda``):
    ``mu`` moves every step by accept/reject, ``lambda`` only when a callback
    replaces ``lm_state.ridge``. The whitened normal matrix
    ``J~'J~ + ridge E`` carries a clean spectral floor at ``ridge`` on the
    metric block regardless of ``W``'s conditioning -- the reason the default
    ``Cholesky()`` path stays accurate at deep ridge.

    ``ridge`` is STRICTLY POSITIVE by contract -- the constructor validates
    it, callbacks must keep it positive, and continuation floors are positive
    -- so ``J~'J~ + ridge E`` is positive definite under the identification
    condition at every reachable state. ``ridge=None`` resolves at ``init``
    to ``sqrt(finfo(dtype).eps)`` of the residual dtype (float64 ~ 1.5e-8,
    float32 ~ 3.5e-4). The metric's factorization never involves the ridge
    weight, so continuation composes unchanged.

    ``linear_solver`` picks the algebra for
    ``(J~'J~ + ridge E + damping I) delta_y = -g`` through a typed config
    (:class:`~nlls_gram.Cholesky` /
    :class:`~nlls_gram.QR` / :class:`~nlls_gram.CG` -- each method's
    knobs live on its own config); every path shares its factorization or
    inner-solve setup between the velocity and geodesic-acceleration solves:

    - ``Cholesky()`` (the default): dense normal equations. ``J'`` is
      materialized per ``jacobian_mode``, ``J~'`` follows by one batched
      ``factor_solve_transpose``, and the assembled
      ``G = J~'J~ + ridge E`` (``ridge`` added on the metric-block diagonal
      only) is cached across rejected steps (only the damping shift
      re-factors).
    - ``QR()``: MINPACK-structured damping-row QR (More 1978), stable at
      small ``ridge``/``damping`` where forming ``G`` squares the condition
      number. One QR of the augmented stack ``[J~, sqrt(ridge) [I 0] | b~]``
      with ``b~ = [r; sqrt(ridge) y_m]`` is cached per ``(x, ridge)`` (the
      extra column carries the ``Q``-transformed residual); each update
      re-factors ``[R; sqrt(damping) I]`` with its ``Q`` retained and solves
      the velocity as a backward-stable least-squares problem at ``cond(A)``,
      never ``cond(A)^2``. The geodesic-acceleration RHS reuses the damped
      factor through corrected semi-normal equations with one fixed
      iterative-refinement pass (Bjorck 1987; Bjorck 1996 Sec. 6.6.5) --
      the second-order correction tolerates the squared conditioning.
    - ``CG(preconditioner, ...)``: matrix-free preconditioned CG on the
      damped normal operator ``J~'J~ + ridge E + damping I`` itself -- the
      same SPD system the ``Cholesky()`` path factors, with products through
      ``jvp``/``vjp`` and the metric's factor callbacks instead of an
      assembled ``G``. The config requires a typed
      :class:`~nlls_gram.Preconditioner`
      (:class:`~nlls_gram.IdentityPreconditioner` opts out): its
      ``apply(v, damping, ctx)`` -- an SPD approximation of the damped
      inverse, handed the same :class:`~nlls_gram.SolverContext` as the
      metric ops -- sits in CG's ``M`` slot with the live damping. The
      operator carries the ``ridge`` spectral floor on the metric block, so
      the preconditioner only has to capture ``J~'J~``'s structure.

    Stopping (``solve``): a ridge solve has two phases -- the residual drops
    to its floor fast, then the iterate slides along the interpolation set
    resolving the null-space (selection) component while ``||r||`` stays
    essentially constant. A pure-residual test is blind to phase 2, so
    ``gtol`` bounds the whitened ridge stationarity ``info.grad_norm =
    ||J~'r + ridge [y_m; 0]||`` (the dual W^{-1}-norm of the half-gradient)
    and ``xtol`` the accepted whitened step norm ``||delta_y||`` (the W-norm
    of the step) -- they mean "done with the current fixed-ridge problem" --
    while ``atol`` is a CONJUNCTIVE filter on the TRUE residual: convergence
    requires (``gtol`` or ``xtol`` fired) AND (``atol == 0`` or
    ``sqrt(resid_loss) <= atol``). atol alone never stops the solve -- a
    pure-residual test would stop at step 0 from any interpolating start
    before the seminorm is minimized -- so ``solve`` rejects ``atol > 0``
    without a positive ``gtol`` or ``xtol``. Calibrating ``gtol`` is clean in
    the whitened geometry: ``info.penalty_grad_norm = sqrt(q(x))``, so
    ``gtol ~ 1e-3 * ridge * sqrt(q(x*))`` resolves the selection to ~1e-3
    relative accuracy, with ``sqrt(q(x*))`` (the solution's seminorm) usually
    known to an order of magnitude before any pilot run. In a continuation
    run the callback keeps lowering ``ridge``, so intermediate stationarity
    with a large residual correctly does not stop the loop.

    ``solve(...).x`` has a custom implicit AD rule with respect to ``p``:
    Gauss-Newton implicit differentiation of the ridge stationarity, posed on
    the whitened variable -- ``(J~'J~ + ridge E) y_dot = -J~'(dr/dp) p_dot``
    then ``x_dot = F_bar^{-1} y_dot`` -- with ``ridge`` frozen
    (stop-gradient) at the returned state's value; the continuation
    schedule's and multi-start selection's dependence on ``p`` is
    deliberately ignored, and no damping enters the AD matrix. Exact
    differentiation carries two extra terms (``sum_i r_i * d2r_i/dx2`` in the
    matrix and ``(dJ'/dp) r`` on the right); both are exactly zero for an
    affine residual and first order in ``||r||`` otherwise, so the rule is
    exact when the converged TRUE residual vanishes. The caveat to state
    plainly: the first-order-in-``||r||`` absolute error translates to a
    small RELATIVE tangent error only under conditioning assumptions
    (``J'J`` bounded below). On a genuinely underdetermined CURVED system the
    null-space block of the exact tangent carries the constraint-curvature
    term ``sum_i nu_i d2r_i`` with multiplier ``nu = r/ridge`` -- which does
    not vanish relative to the ``ridge``-scaled penalty as ``ridge -> 0`` --
    so tangent components read off null-space directions retain an
    O(1)-in-``ridge`` Gauss-Newton bias proportional to the residual
    curvature. ``ad_solver=Cholesky()`` assembles and factors;
    :class:`~nlls_gram.CG` runs matrix-free CG on the same operator with its
    ``preconditioner`` hook (PD does not mean well-conditioned --
    unpreconditioned CG degrades as ``ridge`` shrinks). ``ad_solver=None``
    (the default) matches the forward path: ``Cholesky()`` for the dense
    forwards, and a ``CG`` forward keeps its family AND its preconditioner
    -- the typed apply is damping-analytic, so the forward hook serves the
    undamped system at zero damping exactly (hooks marked
    ``requires_positive_damping`` fall back to unpreconditioned), while the
    AD tolerance and iteration budget stay at the AD defaults; pass
    ``ad_solver=CG(...)`` to pin those. Failed statuses
    return exact zero tangents for ``result.x``/``result.aux`` and evaluate
    the masked tangent program at stop-gradient copies of the caller's
    original inputs and the INITIAL ridge (never a possibly-invalid
    callback-produced value).

    The init/update/solve protocol, callback contract
    (:class:`~nlls_gram.LMContext` ->
    :class:`~nlls_gram.LMAction`), ``multi_start``, ``save_steps``, and
    :class:`~nlls_gram.LMSolveResult` are shared with
    :class:`~nlls_gram.LevenbergMarquardt`; code written against that solver
    ports by changing the constructor (its damping metric -> this metric)
    and reading ``info.resid_loss`` where it means equation error, since
    ``info.loss`` here includes the penalty. Multi-start ranking uses the
    ridge objective at each lane's own final ridge -- comparable across lanes
    when they share a continuation schedule.
    """

    # The metric defines the ridge objective, so a callback-replaced metric
    # is a problem change (convergence suppressed for that step).
    _metric_defines_objective = True

    def __init__(
        self,
        residual_fn,
        *,
        metric,
        ridge=None,
        init_damping=1e-3,
        damping_decrease=0.5,
        damping_increase=4.0,
        linear_solver=Cholesky(),  # noqa: B008 -- frozen, immutable default
        jacobian_mode="auto",
        ad_solver=None,
        has_aux=False,
        cache_jacobian=True,
        geodesic_acceleration=True,
        geodesic_acceptance_ratio=0.75,
    ):
        canonical_residual, residual_arity = canonicalize_residual(residual_fn)
        if (
            ridge is not None
            and not isinstance(ridge, (jax.Array, jax.core.Tracer))
            and float(ridge) <= 0.0
        ):
            raise ValueError(
                "ridge must be strictly positive (ridge = 0 is unsupported: "
                "use AnnealRidge with a positive ridge_floor to "
                "approach the ridgeless limit)"
            )
        if init_damping <= 0 or damping_decrease <= 0 or damping_increase <= 0:
            raise ValueError(
                "init_damping, damping_decrease, and damping_increase must be positive"
            )
        self.residual_fn = canonical_residual
        self.residual_arity = residual_arity
        self.initial_metric = metric
        self._check_registered_instance(metric, "metric")
        self.ridge = ridge
        self.init_damping = init_damping
        self.damping_decrease = damping_decrease
        self.damping_increase = damping_increase
        self.linear_solver = linear_solver
        self.jacobian_mode = jacobian_mode
        self.ad_solver = ad_solver
        self._validate_configuration(linear_solver, ad_solver, penalized=True)
        # The configs' numeric fields fold into the traced LMHyperparams carry
        # at init exactly as the flat constructor args used to; the config
        # instances themselves sit in the value-based static key, so
        # recompile-on-change behavior is unchanged.
        if isinstance(linear_solver, CG):
            # The config requires a typed Preconditioner applying an SPD
            # approximation of (J~'J~ + ridge E + damping I)^{-1}, posed on
            # the whitened variable. IdentityPreconditioner() opts out
            # explicitly.
            if linear_solver.preconditioner is None:
                raise ValueError(
                    "the forward linear_solver requires a preconditioner; "
                    "IdentityPreconditioner() is the explicit opt-out "
                    "(preconditioner=None is legal only in the ad_solver role)"
                )
            self._check_registered_instance(
                linear_solver.preconditioner, "linear_solver.preconditioner"
            )
            self.initial_preconditioner = linear_solver.preconditioner
            # tol=None resolves per residual dtype in hyperparams(), the
            # _ad_cg_tol convention.
            self.iterative_tol = linear_solver.tol
            self.iterative_atol = linear_solver.atol
            self.iterative_maxiter = linear_solver.maxiter
        else:
            self.initial_preconditioner = None
            self.iterative_tol = 0.0
            self.iterative_atol = 0.0
            self.iterative_maxiter = 8
        if isinstance(ad_solver, CG):
            self.ad_solver_tol = ad_solver.tol
            self.ad_solver_atol = ad_solver.atol
            self.ad_solver_maxiter = ad_solver.maxiter
            if ad_solver.preconditioner is None:
                # preconditioner=None in the AD role inherits the CARRIED
                # forward instance at the solution (callback refreshes
                # included) while pinning the AD tolerance and budget.
                self.ad_solver_preconditioner = None
                if self.initial_preconditioner is None:
                    self._ad_preconditioner_source = "none"
                elif self.initial_preconditioner.requires_positive_damping:
                    raise ValueError(
                        "ad_solver preconditioner=None inherits the forward "
                        "preconditioner, but this one divides by the live "
                        "damping and cannot serve the undamped AD system"
                    )
                else:
                    self._ad_preconditioner_source = "carried"
            else:
                if ad_solver.preconditioner.requires_positive_damping:
                    raise ValueError(
                        "this preconditioner divides by the live damping and "
                        "cannot serve in ad_solver (the AD system is undamped)"
                    )
                self._check_registered_instance(
                    ad_solver.preconditioner, "ad_solver.preconditioner"
                )
                self.ad_solver_preconditioner = ad_solver.preconditioner
                self._ad_preconditioner_source = "explicit"
        else:
            self.ad_solver_tol = None
            self.ad_solver_atol = 0.0
            self.ad_solver_maxiter = None
            self.ad_solver_preconditioner = None
            # ad_solver=None matches the forward family, and a CG forward
            # also hands its CARRIED preconditioner to the undamped implicit
            # solve: the AD operator IS the forward operator at zero damping,
            # the typed apply is damping-analytic there, and unpreconditioned
            # implicit CG degrades exactly like the forward as the ridge
            # shrinks. The AD tolerance and budget stay at the AD defaults
            # (run to tolerance); damping-dividing hooks fall back to
            # unpreconditioned.
            inherit = (
                ad_solver is None
                and isinstance(linear_solver, CG)
                and not linear_solver.preconditioner.requires_positive_damping
            )
            self._ad_preconditioner_source = "carried" if inherit else "none"
        self.has_aux = has_aux
        # Only the dense paths materialize J' (and the cholesky/qr caches ride
        # on the same reject-reuse lifecycle), so the flag is inert for the
        # matrix-free normal_cg forward.
        self.cache_jacobian = cache_jacobian and not isinstance(linear_solver, CG)
        self.geodesic_acceleration = geodesic_acceleration
        self.geodesic_acceptance_ratio = geodesic_acceptance_ratio
        # Value-based identity: the jitted solve loop marks the solver itself
        # static, so equal-config solvers built around the same residual and
        # metric-STRUCTURE share the compiled loop across instances. The
        # metric and forward preconditioner key by pytree structure (their
        # arrays are threaded through the carried state); an explicit AD
        # instance keys by identity (its arrays are baked into the tangent
        # program).
        self._static_key = tuple(
            _static_key_component(value)
            for value in (
                residual_fn,
                jax.tree_util.tree_structure(metric),
                ridge,
                init_damping,
                damping_decrease,
                damping_increase,
                _config_static_key(linear_solver, baked=False),
                jacobian_mode,
                _config_static_key(ad_solver, baked=True),
                has_aux,
                self.cache_jacobian,
                geodesic_acceleration,
                geodesic_acceptance_ratio,
            )
        )
        self._static_hash = hash(self._static_key)
        self._sealed = True

    def _resolve_ridge(self, dtype):
        if self.ridge is None:
            return jnp.asarray(jnp.sqrt(jnp.finfo(dtype).eps), dtype=dtype)
        return jnp.asarray(self.ridge, dtype=dtype)

    def init(self, x0, args=None, *, p=None):
        """Build the initial :class:`LMState` at ``x0``.

        One residual evaluation types ``damping`` and resolves
        ``ridge=None`` to the dtype default, and sizes the Jacobian/normal/QR
        cache buffers for the configured dense path. ``hyper`` stays ``None``
        so manual ``update`` loops carry no extra buffers; ``solve``
        populates it for its callbacks.
        """
        self._check_residual_args(args, p)
        residual, aux = self._residual_and_aux(x0, args, p)
        theta, _ = ravel_pytree(x0)
        n_m, _ = self._block_sizes(theta.size)
        dtype = residual.dtype
        damping = jnp.maximum(
            jnp.asarray(self.init_damping, dtype=dtype), _damping_floor(dtype)
        )
        ridge = self._resolve_ridge(dtype)
        instances = dict(
            metric=self.initial_metric, preconditioner=self.initial_preconditioner
        )
        if not self.cache_jacobian:
            return LMState(damping, ridge, **instances)
        p_dim = theta.size
        m = residual.size
        return LMState(
            damping,
            ridge,
            resid=jnp.zeros(residual.shape, dtype=dtype),
            Jt=jnp.zeros((p_dim, m), dtype=dtype),
            jacobian_valid=jnp.asarray(False, dtype=jnp.bool_),
            aux=jax.tree.map(jnp.zeros_like, aux),
            solver_cache=self.linear_solver.new_cache(m, p_dim, n_m, dtype, True),
            **instances,
        )

    def _initial_info(self, x, lm_state, args, p):
        # grad_norm and penalty_grad_norm are +inf sentinels (computing them
        # would cost a Jacobian before the first step) and step_norm is zero;
        # none can satisfy gtol/xtol before any update has run.
        residual, aux = self._residual_and_aux(x, args, p)
        resid_loss = jnp.sum(residual**2)
        theta, _ = ravel_pytree(x)
        ridge = jnp.asarray(lm_state.ridge, dtype=residual.dtype)
        n_m = self._block_sizes(theta.shape[0])[0]
        resolved = self._resolved_state(lm_state)
        ctx = SolverContext(x=theta, lm_state=resolved, args=args, p=p)
        y_m = jnp.asarray(
            resolved.metric.factor_apply(theta[:n_m], ctx), dtype=residual.dtype
        )
        penalty_value = jnp.sum(y_m**2)
        loss = resid_loss + ridge * penalty_value
        zero = jnp.zeros((), dtype=residual.dtype)
        one = jnp.ones((), dtype=residual.dtype)
        infinity = jnp.asarray(jnp.inf, dtype=residual.dtype)
        return LMInfo(
            loss=loss,
            loss_old=loss,
            loss_candidate=loss,
            accepted=jnp.asarray(False, dtype=jnp.bool_),
            damping=jnp.asarray(lm_state.damping, dtype=residual.dtype),
            damping_factor=one,
            used_geodesic=jnp.asarray(False, dtype=jnp.bool_),
            acceleration_ratio=zero,
            grad_norm=infinity,
            step_norm=zero,
            ridge=ridge,
            resid_loss=resid_loss,
            penalty_value=penalty_value,
            penalty_grad_norm=infinity,
            aux=aux,
        )

    def update(self, x, lm_state, args=None, p=None):
        """One LM step on the ridge objective: returns ``(x_new, state, info)``."""
        self._check_residual_args(args, p)
        theta, unravel = ravel_pytree(x)

        if self.has_aux:

            def residual_flat(th):
                value, aux = self.residual_fn(unravel(th), args, p)
                return jnp.ravel(value), aux

            def residual_value(th):
                return residual_flat(th)[0]

        else:

            def residual_flat(th):
                return jnp.ravel(self.residual_fn(unravel(th), args, p))

            residual_value = residual_flat

        # TRUE-residual linearization: matrix-free closures when the linear
        # solver never materializes J, dense J' (reused from the cache after a
        # rejected step) otherwise.
        jvp_fn = JT = Jt = None
        if not self.linear_solver.materializes_jacobian:
            if self.has_aux:
                resid, jvp_fn, aux = jax.linearize(residual_flat, theta, has_aux=True)
            else:
                resid, jvp_fn = jax.linearize(residual_flat, theta)
                aux = None
            transpose_fn = jax.linear_transpose(jvp_fn, theta)

            def JT(cotangent):
                return transpose_fn(cotangent)[0]

        elif self.cache_jacobian:
            resid, Jt, aux = jax.lax.cond(
                lm_state.jacobian_valid,
                lambda _: (lm_state.resid, lm_state.Jt, lm_state.aux),
                lambda _: self._dense_resid_jt_aux(residual_flat, theta),
                operand=None,
            )
        else:
            resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta)

        hyper = (
            lm_state.hyper
            if lm_state.hyper is not None
            else self.hyperparams(resid.dtype)
        )
        damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype)
        damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype)
        damping_floor = _damping_floor(resid.dtype)
        damping = jnp.maximum(
            jnp.asarray(lm_state.damping, dtype=resid.dtype), damping_floor
        )
        if lm_state.ridge is None:
            # A None ridge is a legal LMState (the metric solver leaves it
            # unset); this solver needs one.
            raise ValueError(
                "the lm_state has no ridge; create it with init(x, args, p=p)"
            )
        ridge = jnp.asarray(lm_state.ridge, dtype=resid.dtype)

        # The subproblem is posed on the whitened variable y = F_bar x, where
        # the penalty rows are the constant [I_{n_m} | 0]: the half-gradient is
        # g = F_bar^{-T} J'r + ridge [y_m; 0] (grad F = 2 g; the factor cancels
        # in the LM equations), and every gradient/step quantity below --
        # including the reported norms -- is the whitened one. y_m doubles as
        # the pre-step penalty value ||y_m||^2.
        n_m, n_f = self._block_sizes(theta.shape[0])
        resolved_state = self._resolved_state(lm_state)
        ctx = SolverContext(x=theta, lm_state=resolved_state, args=args, p=p)
        y_m = jnp.asarray(
            resolved_state.metric.factor_apply(theta[:n_m], ctx), dtype=resid.dtype
        )
        penalty_value_old = jnp.sum(y_m**2)
        penalty_gradient = jnp.concatenate([y_m, jnp.zeros(n_f, dtype=resid.dtype)])

        step_solver = self.linear_solver.prepare(
            Subproblem(
                resid=resid,
                theta=theta,
                Jt=Jt,
                jvp_fn=jvp_fn,
                JT=JT,
                whiten=lambda v: self._extended_solve(v, ctx),
                whiten_transpose=lambda v: self._extended_solve_transpose(v, ctx),
                y_m=y_m,
                penalty_gradient=penalty_gradient,
                ridge=ridge,
                damping=damping,
                n_m=n_m,
                n_f=n_f,
                cache=lm_state.solver_cache,
                cache_enabled=self.cache_jacobian,
                hyper=hyper,
                ctx=ctx,
            )
        )
        grad = step_solver.grad

        # First-order step (velocity) and its ridge objective. The solves
        # produce the whitened step delta_y: the x-space step maps back through
        # the factor solve, and the trial penalty uses the linearity of the
        # change of variables -- F_bar(theta + step) = y + delta_y, so no
        # second factor application is ever needed.
        velocity_sub = step_solver.velocity()
        velocity = jnp.asarray(self._extended_solve(velocity_sub, ctx), resid.dtype)

        def trial_penalty(step_sub):
            return jnp.sum((y_m + step_sub[:n_m]) ** 2)

        resid_velocity = residual_value(theta + velocity)
        resid_loss_old = jnp.sum(resid**2)
        loss_old = resid_loss_old + ridge * penalty_value_old
        resid_loss_velocity = jnp.sum(resid_velocity**2)
        penalty_velocity = trial_penalty(velocity_sub)
        loss_velocity = resid_loss_velocity + ridge * penalty_velocity
        zero = jnp.zeros((), dtype=resid.dtype)

        # Geodesic second-order correction, sharing the factorization. The
        # penalty rows are affine, so their directional second derivative is
        # identically zero: f_vv comes from the TRUE residual only and the
        # correction RHS is J' f_vv.
        if self.geodesic_acceleration:
            geodesic_acceptance_ratio = jnp.asarray(
                hyper.geodesic_acceptance_ratio, dtype=resid.dtype
            )

            def first_jvp(th):
                # [1] is the tangent with and without has_aux.
                return jax.jvp(residual_flat, (th,), (velocity,), has_aux=self.has_aux)[
                    1
                ]

            f_vv = jax.jvp(first_jvp, (theta,), (velocity,))[1]
            acceleration_sub = step_solver.correction(f_vv)
            acceleration = jnp.asarray(
                self._extended_solve(acceleration_sub, ctx), dtype=resid.dtype
            )
            accelerated_step = velocity + 0.5 * acceleration
            accelerated_step_sub = velocity_sub + 0.5 * acceleration_sub
            # The ratio criterion lives in the damping geometry's norm -- the
            # whitened (W-norm) one, matching the metric LM's metric_norm.
            acceleration_ratio = (
                2.0
                * jnp.linalg.norm(acceleration_sub)
                / (jnp.linalg.norm(velocity_sub) + jnp.finfo(resid.dtype).eps)
            )
            ratio_accepted = (
                (geodesic_acceptance_ratio > zero)
                & (acceleration_ratio > zero)
                & (acceleration_ratio <= geodesic_acceptance_ratio)
            )

            def accelerated_objective(_):
                resid_accelerated = residual_value(theta + accelerated_step)
                accel_resid_loss = jnp.sum(resid_accelerated**2)
                accel_penalty = trial_penalty(accelerated_step_sub)
                return accel_resid_loss, accel_penalty

            inf = jnp.asarray(jnp.inf, dtype=resid.dtype)
            resid_loss_accelerated, penalty_accelerated = jax.lax.cond(
                ratio_accepted,
                accelerated_objective,
                lambda _: (inf, inf),
                operand=None,
            )
            loss_accelerated = resid_loss_accelerated + ridge * penalty_accelerated
            used_geodesic = ratio_accepted & (loss_accelerated <= loss_velocity)
            step = jnp.where(used_geodesic, accelerated_step, velocity)
            step_sub = jnp.where(used_geodesic, accelerated_step_sub, velocity_sub)
            loss_candidate = jnp.where(used_geodesic, loss_accelerated, loss_velocity)
            resid_loss_candidate = jnp.where(
                used_geodesic, resid_loss_accelerated, resid_loss_velocity
            )
            penalty_candidate = jnp.where(
                used_geodesic, penalty_accelerated, penalty_velocity
            )
        else:
            step = velocity
            step_sub = velocity_sub
            loss_candidate = loss_velocity
            resid_loss_candidate = resid_loss_velocity
            penalty_candidate = penalty_velocity
            used_geodesic = jnp.asarray(False)
            acceleration_ratio = zero

        # Accept iff the ridge objective strictly decreases and is finite --
        # computed with the SAME ridge on both sides, so per-step monotonicity
        # is well-defined even when a callback anneals ridge between steps.
        improved = jnp.isfinite(loss_candidate) & (loss_candidate < loss_old)
        theta_new = jnp.where(improved, theta + step, theta)
        damping_factor = jnp.where(improved, damping_decrease, damping_increase)
        new_damping = jnp.maximum(damping * damping_factor, damping_floor)
        loss = jnp.where(improved, loss_candidate, loss_old)
        resid_loss = jnp.where(improved, resid_loss_candidate, resid_loss_old)
        penalty_value = jnp.where(improved, penalty_candidate, penalty_value_old)
        # Thread the caches built at this step's pre-step (x, ridge):
        # valid = ~improved marks them reusable exactly when the step was
        # rejected (x did not move). ridge and the carried instances pass
        # through unchanged -- only init() and callbacks set them. The input
        # hyper (not the fallback) passes through so the loop carry structure
        # is stable.
        instances = dict(metric=lm_state.metric, preconditioner=lm_state.preconditioner)
        if self.cache_jacobian:
            new_lm_state = LMState(
                new_damping,
                ridge,
                resid,
                Jt,
                ~improved,
                aux,
                lm_state.hyper,
                solver_cache=step_solver.make_cache(~improved),
                **instances,
            )
        else:
            new_lm_state = LMState(
                new_damping, ridge, hyper=lm_state.hyper, **instances
            )
        return (
            unravel(theta_new),
            new_lm_state,
            LMInfo(
                loss=loss,
                loss_old=loss_old,
                loss_candidate=loss_candidate,
                accepted=improved,
                damping=new_damping,
                damping_factor=damping_factor,
                used_geodesic=used_geodesic,
                acceleration_ratio=acceleration_ratio,
                grad_norm=jnp.linalg.norm(grad),
                step_norm=jnp.linalg.norm(step_sub),
                ridge=ridge,
                resid_loss=resid_loss,
                penalty_value=penalty_value,
                penalty_grad_norm=jnp.linalg.norm(penalty_gradient),
                aux=aux,
            ),
        )

    def _validate_tolerances(self, atol, gtol, xtol):
        # atol is a CONJUNCTIVE filter on the true residual, never a stopping
        # rule alone: a residual-only test would stop at any interpolating
        # iterate, before the seminorm is minimized.
        concrete = [not isinstance(t, jax.core.Tracer) for t in (atol, gtol, xtol)]
        if (
            concrete[0]
            and atol > 0
            and (concrete[1] and gtol == 0)
            and (concrete[2] and xtol == 0)
        ):
            raise ValueError(
                "atol > 0 requires a positive gtol or xtol: atol is a "
                "conjunctive filter on the TRUE residual, never a stopping "
                "rule by itself. Calibrate gtol from a pilot run as roughly "
                "1e-3 * ridge * info.penalty_grad_norm"
            )

    def _solve_lm_state(self, x0, args, p, lm_state):
        if lm_state is None:
            # Unconditional init (no minimal-state fast path): resolving
            # ridge=None needs the residual dtype, and the dense caches need
            # their shapes.
            return self.init(x0, args, p=p)
        if lm_state.ridge is None:
            raise ValueError(
                "the caller-supplied lm_state has no ridge; create it with "
                "init(x, args, p=p) or set a positive ridge"
            )
        if (
            not isinstance(lm_state.ridge, jax.core.Tracer)
            and jnp.ndim(lm_state.ridge) == 0
            and float(lm_state.ridge) <= 0.0
        ):
            raise ValueError(
                "the caller-supplied lm_state.ridge must be strictly positive "
                "(ridge = 0 is unsupported)"
            )
        # Recast a hand-replaced ridge to the carried scalar dtype: a
        # weak-typed replace(state, ridge=1e-4) would change the jit input
        # aval and retrace the loop.
        return self._resolved_state(
            dataclasses.replace(
                lm_state,
                ridge=jnp.asarray(
                    lm_state.ridge, dtype=jnp.asarray(lm_state.damping).dtype
                ),
            )
        )

    def _initial_ad_point(self, x, lm_state, args, p):
        # The pre-loop ridge and instances ride along: a failed lane's
        # callback may have left an invalid ridge or metric behind, so the
        # failed tangent uses these.
        return (x, args, p, lm_state.metric, lm_state.preconditioner, lm_state.ridge)

    def _check_action_state(self, lm_state):
        if lm_state.ridge is None:
            raise ValueError(
                "the callback action returned an lm_state without ridge; "
                "use dataclasses.replace(ctx.lm_state, ...) to preserve it"
            )

    def _apply_action_state(self, lm_state, previous):
        # A ridge change leaves the Jacobian cache VALID (J does not depend on
        # ridge) but invalidates the ridge-keyed normal/QR caches and
        # suppresses the convergence check, whose diagnostics were computed at
        # the old ridge. The recast keeps a weak-typed callback float from
        # changing the while_loop carry aval.
        new_ridge = jnp.asarray(lm_state.ridge, dtype=previous.ridge.dtype)
        changed = ~jnp.array_equal(new_ridge, previous.ridge, equal_nan=True)
        return dataclasses.replace(lm_state, ridge=new_ridge), changed

    def _converged(self, info, atol, gtol, xtol):
        # gtol/xtol mean "done with the current fixed-ridge problem"; atol is
        # a CONJUNCTIVE filter on the TRUE residual, never sufficient alone.
        gtol_met = (gtol > 0) & (info.grad_norm < gtol)
        xtol_met = (xtol > 0) & info.accepted & (info.step_norm < xtol)
        residual_ok = (atol <= 0) | (jnp.sqrt(info.resid_loss) <= atol)
        return (gtol_met | xtol_met) & residual_ok

    def _cast_state(self, lm_state, dtype):
        updates = dict(
            damping=jnp.asarray(lm_state.damping, dtype=dtype),
            ridge=jnp.asarray(lm_state.ridge, dtype=dtype),
            hyper=_cast_hyper(lm_state.hyper, dtype),
        )
        if lm_state.solver_cache is not None:
            cache = lm_state.solver_cache
            updates["solver_cache"] = dataclasses.replace(
                cache, ridge=jnp.asarray(cache.ridge, dtype=dtype)
            )
        return dataclasses.replace(lm_state, **updates)

    def _ranking_objective(self, result, p, callback):
        # Multi-start selection ranks by the ridge objective at each lane's
        # OWN final ridge (comparable across lanes when they share a
        # continuation schedule). Without a callback info.loss already reports
        # it at the retained iterate; a callback can replace x/args/ridge
        # after the last update, so recompute. Nonfinite masks to +inf.
        if callback is None:
            loss = result.info.loss
        else:
            residual = self._residual_and_aux(result.x, result.args, p)[0]
            theta, _ = ravel_pytree(result.x)
            n_m = self._block_sizes(theta.shape[0])[0]
            resolved = self._resolved_state(result.lm_state)
            ctx = SolverContext(x=theta, lm_state=resolved, args=result.args, p=p)
            ridge = jnp.asarray(result.lm_state.ridge, dtype=residual.dtype)
            y_m = jnp.asarray(
                resolved.metric.factor_apply(theta[:n_m], ctx), dtype=residual.dtype
            )
            loss = jnp.sum(residual**2) + ridge * jnp.sum(y_m**2)
        return jnp.where(
            jnp.isfinite(loss), loss, jnp.asarray(jnp.inf, dtype=loss.dtype)
        )

    def _resolved_ad_solver(self):
        if self.ad_solver is None:
            # Matrix-free forward -> matrix-free AD.
            return "normal_cg" if isinstance(self.linear_solver, CG) else "cholesky"
        return "cholesky" if isinstance(self.ad_solver, Cholesky) else "normal_cg"

    def _ad_x_tangent(self, x, args, p, p_dot, result, ad_success, initial_ad_point):
        if p is None:
            return jax.tree.map(_zero_tangent_leaf, x)
        # A successful tangent uses the winner's own final ridge and carried
        # instances; a failed one the pre-loop initial ones (a callback may
        # have left invalid values behind). All are stop-gradient'd -- inert
        # conditioning data for the factor callbacks.
        final_ridge = jax.lax.stop_gradient(result.lm_state.ridge)
        initial_ridge = jax.lax.stop_gradient(initial_ad_point[5])
        ridge = jnp.where(
            ad_success, final_ridge, jnp.asarray(initial_ridge, final_ridge.dtype)
        )
        lm_state = jax.lax.stop_gradient(result.lm_state)
        initial_instances = jax.lax.stop_gradient(initial_ad_point[3:5])
        lm_state = dataclasses.replace(
            lm_state,
            metric=_where_tree(ad_success, lm_state.metric, initial_instances[0]),
            preconditioner=_where_tree(
                ad_success, lm_state.preconditioner, initial_instances[1]
            ),
        )
        if self._resolved_ad_solver() == "cholesky":
            return self._ad_tangent_cholesky(x, args, p, p_dot, ridge, lm_state)
        return self._ad_tangent_normal_cg(x, args, p, p_dot, ridge, lm_state)

    def _ad_tangent_cholesky(self, x, args, p, p_dot, ridge, lm_state):
        # The GN implicit rule posed on the whitened variable y = F_bar x:
        # (J~'J~ + ridge E) y_dot = -J~'(dr/dp) p_dot, then
        # x_dot = F_bar^{-1} y_dot -- no damping; the matrix is PD under the
        # identification condition because ridge > 0 by contract.
        theta, unravel, residual, theta_jvp, residual_p_dot = self._ad_linearization(
            x, args, p, p_dot
        )
        n_m = self._block_sizes(theta.shape[0])[0]
        ctx = SolverContext(x=theta, lm_state=lm_state, args=args, p=p)
        Jt = self._assemble_jt(theta_jvp, theta, residual)
        ridge_typed = jnp.asarray(ridge, dtype=residual.dtype)
        Jt_sub = jnp.asarray(
            self._extended_solve_transpose(Jt, ctx), dtype=residual.dtype
        )
        diag = jnp.arange(n_m)
        normal_matrix = mm(Jt_sub, Jt_sub.T).at[diag, diag].add(ridge_typed)
        factor = jsp_linalg.cho_factor(normal_matrix)
        y_dot = jsp_linalg.cho_solve(factor, -mm(Jt_sub, residual_p_dot))
        theta_dot = jnp.asarray(self._extended_solve(y_dot, ctx), residual.dtype)
        return unravel(theta_dot)

    def _ad_tangent_normal_cg(self, x, args, p, p_dot, ridge, lm_state):
        # Matrix-free CG on the same whitened PD operator (see
        # _ad_tangent_cholesky); matvec = J~'(J~ u) + ridge [u_m; 0].
        theta, unravel, residual, theta_jvp, residual_p_dot = self._ad_linearization(
            x, args, p, p_dot
        )
        n_m, n_f = self._block_sizes(theta.shape[0])
        ctx = SolverContext(x=theta, lm_state=lm_state, args=args, p=p)
        theta_transpose = jax.linear_transpose(theta_jvp, theta)

        def JT(cotangent):
            return theta_transpose(cotangent)[0]

        ridge_typed = jnp.asarray(ridge, dtype=residual.dtype)

        def normal_matvec(u):
            gauss_newton = jnp.asarray(
                self._extended_solve_transpose(
                    JT(
                        theta_jvp(
                            jnp.asarray(self._extended_solve(u, ctx), residual.dtype)
                        )
                    ),
                    ctx,
                ),
                dtype=residual.dtype,
            )
            metric_shift = jnp.concatenate(
                [u[:n_m], jnp.zeros(n_f, dtype=residual.dtype)]
            )
            return gauss_newton + ridge_typed * metric_shift

        cg_tol = self._ad_cg_tol(residual.dtype)
        cg_atol = jnp.asarray(self.ad_solver_atol, dtype=residual.dtype)
        ad_preconditioner = self._ad_preconditioner(lm_state)
        if ad_preconditioner is None:
            apply_M = None
        else:
            # The AD system is undamped, so the preconditioner sees zero
            # damping (requires_positive_damping subclasses were rejected at
            # construction).
            zero_damping = jnp.zeros((), dtype=residual.dtype)

            def apply_M(v):
                return ad_preconditioner.apply(v, zero_damping, ctx)

        def solve(matvec, rhs_value):
            solution, _ = jsp_sparse_linalg.cg(
                matvec,
                rhs_value,
                tol=cg_tol,
                atol=cg_atol,
                maxiter=self.ad_solver_maxiter,
                M=apply_M,
            )
            return solution

        rhs = -jnp.asarray(
            self._extended_solve_transpose(JT(residual_p_dot), ctx),
            dtype=residual.dtype,
        )
        y_dot = jax.lax.custom_linear_solve(
            normal_matvec,
            rhs,
            solve,
            symmetric=True,
        )
        theta_dot = jnp.asarray(self._extended_solve(y_dot, ctx), residual.dtype)
        return unravel(theta_dot)

init(x0, args=None, *, p=None)

Build the initial :class:LMState at x0.

One residual evaluation types damping and resolves ridge=None to the dtype default, and sizes the Jacobian/normal/QR cache buffers for the configured dense path. hyper stays None so manual update loops carry no extra buffers; solve populates it for its callbacks.

Source code in src/nlls_gram/ridge_lm.py

def init(self, x0, args=None, *, p=None):
    """Build the initial :class:`LMState` at ``x0``.

    One residual evaluation types ``damping`` and resolves
    ``ridge=None`` to the dtype default, and sizes the Jacobian/normal/QR
    cache buffers for the configured dense path. ``hyper`` stays ``None``
    so manual ``update`` loops carry no extra buffers; ``solve``
    populates it for its callbacks.
    """
    self._check_residual_args(args, p)
    residual, aux = self._residual_and_aux(x0, args, p)
    theta, _ = ravel_pytree(x0)
    n_m, _ = self._block_sizes(theta.size)
    dtype = residual.dtype
    damping = jnp.maximum(
        jnp.asarray(self.init_damping, dtype=dtype), _damping_floor(dtype)
    )
    ridge = self._resolve_ridge(dtype)
    instances = dict(
        metric=self.initial_metric, preconditioner=self.initial_preconditioner
    )
    if not self.cache_jacobian:
        return LMState(damping, ridge, **instances)
    p_dim = theta.size
    m = residual.size
    return LMState(
        damping,
        ridge,
        resid=jnp.zeros(residual.shape, dtype=dtype),
        Jt=jnp.zeros((p_dim, m), dtype=dtype),
        jacobian_valid=jnp.asarray(False, dtype=jnp.bool_),
        aux=jax.tree.map(jnp.zeros_like, aux),
        solver_cache=self.linear_solver.new_cache(m, p_dim, n_m, dtype, True),
        **instances,
    )

update(x, lm_state, args=None, p=None)

One LM step on the ridge objective: returns (x_new, state, info).

Source code in src/nlls_gram/ridge_lm.py

def update(self, x, lm_state, args=None, p=None):
    """One LM step on the ridge objective: returns ``(x_new, state, info)``."""
    self._check_residual_args(args, p)
    theta, unravel = ravel_pytree(x)

    if self.has_aux:

        def residual_flat(th):
            value, aux = self.residual_fn(unravel(th), args, p)
            return jnp.ravel(value), aux

        def residual_value(th):
            return residual_flat(th)[0]

    else:

        def residual_flat(th):
            return jnp.ravel(self.residual_fn(unravel(th), args, p))

        residual_value = residual_flat

    # TRUE-residual linearization: matrix-free closures when the linear
    # solver never materializes J, dense J' (reused from the cache after a
    # rejected step) otherwise.
    jvp_fn = JT = Jt = None
    if not self.linear_solver.materializes_jacobian:
        if self.has_aux:
            resid, jvp_fn, aux = jax.linearize(residual_flat, theta, has_aux=True)
        else:
            resid, jvp_fn = jax.linearize(residual_flat, theta)
            aux = None
        transpose_fn = jax.linear_transpose(jvp_fn, theta)

        def JT(cotangent):
            return transpose_fn(cotangent)[0]

    elif self.cache_jacobian:
        resid, Jt, aux = jax.lax.cond(
            lm_state.jacobian_valid,
            lambda _: (lm_state.resid, lm_state.Jt, lm_state.aux),
            lambda _: self._dense_resid_jt_aux(residual_flat, theta),
            operand=None,
        )
    else:
        resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta)

    hyper = (
        lm_state.hyper
        if lm_state.hyper is not None
        else self.hyperparams(resid.dtype)
    )
    damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype)
    damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype)
    damping_floor = _damping_floor(resid.dtype)
    damping = jnp.maximum(
        jnp.asarray(lm_state.damping, dtype=resid.dtype), damping_floor
    )
    if lm_state.ridge is None:
        # A None ridge is a legal LMState (the metric solver leaves it
        # unset); this solver needs one.
        raise ValueError(
            "the lm_state has no ridge; create it with init(x, args, p=p)"
        )
    ridge = jnp.asarray(lm_state.ridge, dtype=resid.dtype)

    # The subproblem is posed on the whitened variable y = F_bar x, where
    # the penalty rows are the constant [I_{n_m} | 0]: the half-gradient is
    # g = F_bar^{-T} J'r + ridge [y_m; 0] (grad F = 2 g; the factor cancels
    # in the LM equations), and every gradient/step quantity below --
    # including the reported norms -- is the whitened one. y_m doubles as
    # the pre-step penalty value ||y_m||^2.
    n_m, n_f = self._block_sizes(theta.shape[0])
    resolved_state = self._resolved_state(lm_state)
    ctx = SolverContext(x=theta, lm_state=resolved_state, args=args, p=p)
    y_m = jnp.asarray(
        resolved_state.metric.factor_apply(theta[:n_m], ctx), dtype=resid.dtype
    )
    penalty_value_old = jnp.sum(y_m**2)
    penalty_gradient = jnp.concatenate([y_m, jnp.zeros(n_f, dtype=resid.dtype)])

    step_solver = self.linear_solver.prepare(
        Subproblem(
            resid=resid,
            theta=theta,
            Jt=Jt,
            jvp_fn=jvp_fn,
            JT=JT,
            whiten=lambda v: self._extended_solve(v, ctx),
            whiten_transpose=lambda v: self._extended_solve_transpose(v, ctx),
            y_m=y_m,
            penalty_gradient=penalty_gradient,
            ridge=ridge,
            damping=damping,
            n_m=n_m,
            n_f=n_f,
            cache=lm_state.solver_cache,
            cache_enabled=self.cache_jacobian,
            hyper=hyper,
            ctx=ctx,
        )
    )
    grad = step_solver.grad

    # First-order step (velocity) and its ridge objective. The solves
    # produce the whitened step delta_y: the x-space step maps back through
    # the factor solve, and the trial penalty uses the linearity of the
    # change of variables -- F_bar(theta + step) = y + delta_y, so no
    # second factor application is ever needed.
    velocity_sub = step_solver.velocity()
    velocity = jnp.asarray(self._extended_solve(velocity_sub, ctx), resid.dtype)

    def trial_penalty(step_sub):
        return jnp.sum((y_m + step_sub[:n_m]) ** 2)

    resid_velocity = residual_value(theta + velocity)
    resid_loss_old = jnp.sum(resid**2)
    loss_old = resid_loss_old + ridge * penalty_value_old
    resid_loss_velocity = jnp.sum(resid_velocity**2)
    penalty_velocity = trial_penalty(velocity_sub)
    loss_velocity = resid_loss_velocity + ridge * penalty_velocity
    zero = jnp.zeros((), dtype=resid.dtype)

    # Geodesic second-order correction, sharing the factorization. The
    # penalty rows are affine, so their directional second derivative is
    # identically zero: f_vv comes from the TRUE residual only and the
    # correction RHS is J' f_vv.
    if self.geodesic_acceleration:
        geodesic_acceptance_ratio = jnp.asarray(
            hyper.geodesic_acceptance_ratio, dtype=resid.dtype
        )

        def first_jvp(th):
            # [1] is the tangent with and without has_aux.
            return jax.jvp(residual_flat, (th,), (velocity,), has_aux=self.has_aux)[
                1
            ]

        f_vv = jax.jvp(first_jvp, (theta,), (velocity,))[1]
        acceleration_sub = step_solver.correction(f_vv)
        acceleration = jnp.asarray(
            self._extended_solve(acceleration_sub, ctx), dtype=resid.dtype
        )
        accelerated_step = velocity + 0.5 * acceleration
        accelerated_step_sub = velocity_sub + 0.5 * acceleration_sub
        # The ratio criterion lives in the damping geometry's norm -- the
        # whitened (W-norm) one, matching the metric LM's metric_norm.
        acceleration_ratio = (
            2.0
            * jnp.linalg.norm(acceleration_sub)
            / (jnp.linalg.norm(velocity_sub) + jnp.finfo(resid.dtype).eps)
        )
        ratio_accepted = (
            (geodesic_acceptance_ratio > zero)
            & (acceleration_ratio > zero)
            & (acceleration_ratio <= geodesic_acceptance_ratio)
        )

        def accelerated_objective(_):
            resid_accelerated = residual_value(theta + accelerated_step)
            accel_resid_loss = jnp.sum(resid_accelerated**2)
            accel_penalty = trial_penalty(accelerated_step_sub)
            return accel_resid_loss, accel_penalty

        inf = jnp.asarray(jnp.inf, dtype=resid.dtype)
        resid_loss_accelerated, penalty_accelerated = jax.lax.cond(
            ratio_accepted,
            accelerated_objective,
            lambda _: (inf, inf),
            operand=None,
        )
        loss_accelerated = resid_loss_accelerated + ridge * penalty_accelerated
        used_geodesic = ratio_accepted & (loss_accelerated <= loss_velocity)
        step = jnp.where(used_geodesic, accelerated_step, velocity)
        step_sub = jnp.where(used_geodesic, accelerated_step_sub, velocity_sub)
        loss_candidate = jnp.where(used_geodesic, loss_accelerated, loss_velocity)
        resid_loss_candidate = jnp.where(
            used_geodesic, resid_loss_accelerated, resid_loss_velocity
        )
        penalty_candidate = jnp.where(
            used_geodesic, penalty_accelerated, penalty_velocity
        )
    else:
        step = velocity
        step_sub = velocity_sub
        loss_candidate = loss_velocity
        resid_loss_candidate = resid_loss_velocity
        penalty_candidate = penalty_velocity
        used_geodesic = jnp.asarray(False)
        acceleration_ratio = zero

    # Accept iff the ridge objective strictly decreases and is finite --
    # computed with the SAME ridge on both sides, so per-step monotonicity
    # is well-defined even when a callback anneals ridge between steps.
    improved = jnp.isfinite(loss_candidate) & (loss_candidate < loss_old)
    theta_new = jnp.where(improved, theta + step, theta)
    damping_factor = jnp.where(improved, damping_decrease, damping_increase)
    new_damping = jnp.maximum(damping * damping_factor, damping_floor)
    loss = jnp.where(improved, loss_candidate, loss_old)
    resid_loss = jnp.where(improved, resid_loss_candidate, resid_loss_old)
    penalty_value = jnp.where(improved, penalty_candidate, penalty_value_old)
    # Thread the caches built at this step's pre-step (x, ridge):
    # valid = ~improved marks them reusable exactly when the step was
    # rejected (x did not move). ridge and the carried instances pass
    # through unchanged -- only init() and callbacks set them. The input
    # hyper (not the fallback) passes through so the loop carry structure
    # is stable.
    instances = dict(metric=lm_state.metric, preconditioner=lm_state.preconditioner)
    if self.cache_jacobian:
        new_lm_state = LMState(
            new_damping,
            ridge,
            resid,
            Jt,
            ~improved,
            aux,
            lm_state.hyper,
            solver_cache=step_solver.make_cache(~improved),
            **instances,
        )
    else:
        new_lm_state = LMState(
            new_damping, ridge, hyper=lm_state.hyper, **instances
        )
    return (
        unravel(theta_new),
        new_lm_state,
        LMInfo(
            loss=loss,
            loss_old=loss_old,
            loss_candidate=loss_candidate,
            accepted=improved,
            damping=new_damping,
            damping_factor=damping_factor,
            used_geodesic=used_geodesic,
            acceleration_ratio=acceleration_ratio,
            grad_norm=jnp.linalg.norm(grad),
            step_norm=jnp.linalg.norm(step_sub),
            ridge=ridge,
            resid_loss=resid_loss,
            penalty_value=penalty_value,
            penalty_grad_norm=jnp.linalg.norm(penalty_gradient),
            aux=aux,
        ),
    )

nlls_gram.LevenbergMarquardt

Bases: LevenbergMarquardtBase

Levenberg-Marquardt least squares over a JAX pytree x.

residual_fn takes (x), (x, args), or (x, args, p) and returns a residual pytree (or (residual, aux) with has_aux=True). args is solver-inert auxiliary data; p is what solve(...).x carries an implicit derivative with respect to.

metric (default None -- Euclidean) is a :class:~nlls_gram.Metric covering the leading metric.size coordinates of the flattened iterate, with the rest a free block. It defines the damping geometry: the subproblem is min ||r + J s||^2 + damping ||s||_W^2, so the damping -> 0 limit is the minimum-W-norm correction.

linear_solver is a typed config -- :class:~nlls_gram.Cholesky (the default; form="auto" factors the smaller of the m x m dual and the n x n normal system), :class:~nlls_gram.QR (damping-row QR, stable at tiny damping and rank-safe), :class:~nlls_gram.CG (matrix-free in parameter space), or :class:~nlls_gram.GramCG (matrix-free in residual space, the m << n form). ad_solver takes None (match the forward family), Cholesky(), :class:~nlls_gram.SVD (the pseudoinverse rule, for a rank-deficient undamped tangent), CG(...), or GramCG(...).

init/update/solve, the callback protocol, save_steps, multi_start, and implicit AD are shared with the ridge solver; info.loss here is the plain sum of squared residuals. Stopping is disjunctive: atol on the residual norm, gtol on the whitened stationarity info.grad_norm, xtol on an accepted step's whitened norm; any one firing reports CONVERGED.

Source code in src/nlls_gram/metric_lm.py

class LevenbergMarquardt(LevenbergMarquardtBase):
    """Levenberg-Marquardt least squares over a JAX pytree ``x``.

    ``residual_fn`` takes ``(x)``, ``(x, args)``, or ``(x, args, p)`` and
    returns a residual pytree (or ``(residual, aux)`` with ``has_aux=True``).
    ``args`` is solver-inert auxiliary data; ``p`` is what ``solve(...).x``
    carries an implicit derivative with respect to.

    ``metric`` (default ``None`` -- Euclidean) is a
    :class:`~nlls_gram.Metric` covering the leading ``metric.size``
    coordinates of the flattened iterate, with the rest a free block. It
    defines the damping geometry: the subproblem is
    ``min ||r + J s||^2 + damping ||s||_W^2``, so the ``damping -> 0`` limit
    is the minimum-``W``-norm correction.

    ``linear_solver`` is a typed config -- :class:`~nlls_gram.Cholesky` (the
    default; ``form="auto"`` factors the smaller of the ``m x m`` dual and the
    ``n x n`` normal system), :class:`~nlls_gram.QR` (damping-row QR, stable
    at tiny damping and rank-safe), :class:`~nlls_gram.CG` (matrix-free in
    parameter space), or :class:`~nlls_gram.GramCG` (matrix-free in residual
    space, the ``m << n`` form). ``ad_solver`` takes ``None`` (match the
    forward family), ``Cholesky()``, :class:`~nlls_gram.SVD` (the
    pseudoinverse rule, for a rank-deficient undamped tangent), ``CG(...)``,
    or ``GramCG(...)``.

    ``init``/``update``/``solve``, the callback protocol, ``save_steps``,
    ``multi_start``, and implicit AD are shared with the ridge solver;
    ``info.loss`` here is the plain sum of squared residuals. Stopping is
    disjunctive: ``atol`` on the residual norm, ``gtol`` on the whitened
    stationarity ``info.grad_norm``, ``xtol`` on an accepted step's whitened
    norm; any one firing reports ``CONVERGED``.
    """

    def __init__(
        self,
        residual_fn,
        *,
        metric=None,
        init_damping=1e-3,
        damping_decrease=0.5,
        damping_increase=4.0,
        linear_solver=Cholesky(),  # noqa: B008 -- frozen, immutable default
        jacobian_mode="auto",
        ad_solver=None,
        has_aux=False,
        cache_jacobian=True,
        geodesic_acceleration=True,
        geodesic_acceptance_ratio=0.75,
    ):
        canonical_residual, residual_arity = canonicalize_residual(residual_fn)
        if init_damping <= 0 or damping_decrease <= 0 or damping_increase <= 0:
            raise ValueError(
                "init_damping, damping_decrease, and damping_increase must be positive"
            )
        self.residual_fn = canonical_residual
        self.residual_arity = residual_arity
        self.initial_metric = _EuclideanMetric() if metric is None else metric
        self._check_registered_instance(self.initial_metric, "metric")
        self.init_damping = init_damping
        self.damping_decrease = damping_decrease
        self.damping_increase = damping_increase
        self.linear_solver = linear_solver
        self.jacobian_mode = jacobian_mode
        self.ad_solver = ad_solver
        self._validate_configuration(linear_solver, ad_solver, penalized=False)
        krylov = isinstance(linear_solver, (CG, GramCG))
        if krylov:
            if linear_solver.preconditioner is None:
                raise ValueError(
                    "the forward linear_solver requires a preconditioner; "
                    "IdentityPreconditioner() is the explicit opt-out "
                    "(preconditioner=None is legal only in the ad_solver role)"
                )
            self._check_registered_instance(
                linear_solver.preconditioner, "linear_solver.preconditioner"
            )
            self.initial_preconditioner = linear_solver.preconditioner
            self.iterative_tol = linear_solver.tol
            self.iterative_atol = linear_solver.atol
            self.iterative_maxiter = linear_solver.maxiter
        else:
            self.initial_preconditioner = None
            self.iterative_tol = 0.0
            self.iterative_atol = 0.0
            self.iterative_maxiter = 8
        if isinstance(ad_solver, (CG, GramCG)):
            self.ad_solver_tol = ad_solver.tol
            self.ad_solver_atol = ad_solver.atol
            self.ad_solver_maxiter = ad_solver.maxiter
            self.ad_solver_penalty = getattr(ad_solver, "penalty", None)
            if ad_solver.preconditioner is None:
                # preconditioner=None in the AD role inherits the CARRIED
                # forward instance at the solution (callback refreshes
                # included) while pinning the AD tolerance and budget.
                self.ad_solver_preconditioner = None
                if self.initial_preconditioner is None:
                    self._ad_preconditioner_source = "none"
                elif self.initial_preconditioner.requires_positive_damping:
                    raise ValueError(
                        "ad_solver preconditioner=None inherits the forward "
                        "preconditioner, but this one divides by the live "
                        "damping and cannot serve the undamped AD system"
                    )
                else:
                    self._ad_preconditioner_source = "carried"
            else:
                if ad_solver.preconditioner.requires_positive_damping:
                    raise ValueError(
                        "this preconditioner divides by the live damping and "
                        "cannot serve in ad_solver (the AD system is undamped)"
                    )
                self._check_registered_instance(
                    ad_solver.preconditioner, "ad_solver.preconditioner"
                )
                self.ad_solver_preconditioner = ad_solver.preconditioner
                self._ad_preconditioner_source = "explicit"
        else:
            self.ad_solver_tol = None
            self.ad_solver_atol = 0.0
            self.ad_solver_maxiter = None
            self.ad_solver_penalty = None
            self.ad_solver_preconditioner = None
            # ad_solver=None under a matrix-free forward hands the CARRIED
            # forward preconditioner to the undamped implicit solve: the AD
            # operator IS the forward operator at zero damping.
            # Damping-dividing hooks fall back to unpreconditioned.
            inherit = (
                ad_solver is None
                and krylov
                and not linear_solver.preconditioner.requires_positive_damping
            )
            self._ad_preconditioner_source = "carried" if inherit else "none"
            if inherit:
                self.ad_solver_penalty = getattr(linear_solver, "penalty", None)
        self.has_aux = has_aux
        # Only the dense paths materialize J', and the caches ride the same
        # reject-reuse lifecycle, so the flag is inert for the matrix-free forms.
        self.cache_jacobian = cache_jacobian and linear_solver.materializes_jacobian
        self.geodesic_acceleration = geodesic_acceleration
        self.geodesic_acceptance_ratio = geodesic_acceptance_ratio
        # Metric and forward-preconditioner instances key by pytree structure:
        # their arrays are threaded through the carried state, so equal-config
        # fresh instances share one compiled loop. An explicit AD instance is
        # baked into the tangent program as constants, so it keys by identity.
        self._static_key = tuple(
            _static_key_component(value)
            for value in (
                residual_fn,
                jax.tree_util.tree_structure(self.initial_metric),
                init_damping,
                damping_decrease,
                damping_increase,
                _config_static_key(linear_solver, baked=False),
                jacobian_mode,
                _config_static_key(ad_solver, baked=True),
                has_aux,
                self.cache_jacobian,
                geodesic_acceleration,
                geodesic_acceptance_ratio,
            )
        )
        self._static_hash = hash(self._static_key)
        self._sealed = True

    def init(self, x0, args=None, *, p=None):
        """Build the initial :class:`~nlls_gram.LMState` at ``x0``.

        One residual evaluation types ``damping`` and sizes the Jacobian and
        linear-solver cache buffers. ``hyper`` stays ``None`` so manual
        ``update`` loops carry no extra buffers; ``solve`` populates it.
        """
        self._check_residual_args(args, p)
        residual, aux = self._residual_and_aux(x0, args, p)
        theta, _ = ravel_pytree(x0)
        n_m, _ = self._block_sizes(theta.size)
        dtype = residual.dtype
        damping = jnp.maximum(
            jnp.asarray(self.init_damping, dtype=dtype), _damping_floor(dtype)
        )
        instances = dict(
            metric=self.initial_metric, preconditioner=self.initial_preconditioner
        )
        if not self.cache_jacobian:
            return LMState(damping, **instances)
        return LMState(
            damping,
            resid=jnp.zeros(residual.shape, dtype=dtype),
            Jt=jnp.zeros((theta.size, residual.size), dtype=dtype),
            jacobian_valid=jnp.asarray(False, dtype=jnp.bool_),
            aux=jax.tree.map(jnp.zeros_like, aux),
            solver_cache=self.linear_solver.new_cache(
                residual.size, theta.size, n_m, dtype, False
            ),
            **instances,
        )

    def _solve_lm_state(self, x0, args, p, lm_state):
        if lm_state is not None:
            # A hand-built pre-seeding state enters the loop with the
            # constructor instances; the carry needs them present.
            return self._resolved_state(lm_state)
        if self.cache_jacobian:
            return self.init(x0, args, p=p)
        # Nothing needs sizing from a residual evaluation, so skip it: the
        # loop recasts the damping dtype itself.
        return LMState(
            jnp.asarray(self.init_damping),
            metric=self.initial_metric,
            preconditioner=self.initial_preconditioner,
        )

    def _initial_info(self, x, lm_state, args, p):
        # grad_norm is a +inf sentinel (computing it would cost a Jacobian
        # before the first step) and step_norm is zero; neither can satisfy
        # gtol/xtol before any update has run.
        residual, aux = self._residual_and_aux(x, args, p)
        loss = jnp.sum(residual**2)
        zero = jnp.zeros((), dtype=residual.dtype)
        return LMInfo(
            loss=loss,
            loss_old=loss,
            loss_candidate=loss,
            accepted=jnp.asarray(False, dtype=jnp.bool_),
            damping=jnp.asarray(lm_state.damping, dtype=residual.dtype),
            damping_factor=jnp.ones((), dtype=residual.dtype),
            used_geodesic=jnp.asarray(False, dtype=jnp.bool_),
            acceleration_ratio=zero,
            grad_norm=jnp.asarray(jnp.inf, dtype=residual.dtype),
            step_norm=zero,
            aux=aux,
        )

    def update(self, x, lm_state, args=None, p=None):
        """One LM step: returns ``(x_new, lm_state, info)``."""
        self._check_residual_args(args, p)
        theta, unravel = ravel_pytree(x)

        if self.has_aux:

            def residual_flat(th):
                value, aux = self.residual_fn(unravel(th), args, p)
                return jnp.ravel(value), aux

            def residual_value(th):
                return residual_flat(th)[0]

        else:

            def residual_flat(th):
                return jnp.ravel(self.residual_fn(unravel(th), args, p))

            residual_value = residual_flat

        jvp_fn = JT = Jt = None
        if not self.linear_solver.materializes_jacobian:
            if self.has_aux:
                resid, jvp_fn, aux = jax.linearize(residual_flat, theta, has_aux=True)
            else:
                resid, jvp_fn = jax.linearize(residual_flat, theta)
                aux = None
            transpose_fn = jax.linear_transpose(jvp_fn, theta)

            def JT(cotangent):
                return transpose_fn(cotangent)[0]

        elif self.cache_jacobian:
            resid, Jt, aux = jax.lax.cond(
                lm_state.jacobian_valid,
                lambda _: (lm_state.resid, lm_state.Jt, lm_state.aux),
                lambda _: self._dense_resid_jt_aux(residual_flat, theta),
                operand=None,
            )
        else:
            resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta)

        hyper = (
            lm_state.hyper
            if lm_state.hyper is not None
            else self.hyperparams(resid.dtype)
        )
        damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype)
        damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype)
        damping_floor = _damping_floor(resid.dtype)
        damping = jnp.maximum(
            jnp.asarray(lm_state.damping, dtype=resid.dtype), damping_floor
        )

        n_m, n_f = self._block_sizes(theta.shape[0])
        ctx = SolverContext(
            x=theta, lm_state=self._resolved_state(lm_state), args=args, p=p
        )
        zero = jnp.zeros((), dtype=resid.dtype)
        step_solver = self.linear_solver.prepare(
            Subproblem(
                resid=resid,
                theta=theta,
                Jt=Jt,
                jvp_fn=jvp_fn,
                JT=JT,
                whiten=lambda v: self._extended_solve(v, ctx),
                whiten_transpose=lambda v: self._extended_solve_transpose(v, ctx),
                y_m=jnp.zeros(n_m, dtype=resid.dtype),
                penalty_gradient=jnp.zeros(theta.shape[0], dtype=resid.dtype),
                ridge=zero,
                damping=damping,
                n_m=n_m,
                n_f=n_f,
                cache=lm_state.solver_cache,
                cache_enabled=self.cache_jacobian,
                hyper=hyper,
                ctx=ctx,
                penalized=False,
            )
        )

        # The solves produce the whitened step; the x-space step maps back
        # through the factor solve.
        velocity_sub = step_solver.velocity()
        velocity = jnp.asarray(self._extended_solve(velocity_sub, ctx), resid.dtype)
        loss_old = jnp.sum(resid**2)
        resid_velocity = residual_value(theta + velocity)
        loss_velocity = jnp.sum(resid_velocity**2)

        # Geodesic second-order correction, sharing the factorization.
        if self.geodesic_acceleration:
            geodesic_acceptance_ratio = jnp.asarray(
                hyper.geodesic_acceptance_ratio, dtype=resid.dtype
            )

            def first_jvp(th):
                # [1] is the tangent with and without has_aux.
                return jax.jvp(residual_flat, (th,), (velocity,), has_aux=self.has_aux)[
                    1
                ]

            f_vv = jax.jvp(first_jvp, (theta,), (velocity,))[1]
            acceleration_sub = step_solver.correction(f_vv)
            acceleration = jnp.asarray(
                self._extended_solve(acceleration_sub, ctx), dtype=resid.dtype
            )
            accelerated_step = velocity + 0.5 * acceleration
            # The ratio criterion lives in the damping geometry's norm -- the
            # whitened one.
            acceleration_ratio = (
                2.0
                * jnp.linalg.norm(acceleration_sub)
                / (jnp.linalg.norm(velocity_sub) + jnp.finfo(resid.dtype).eps)
            )
            ratio_accepted = (
                (geodesic_acceptance_ratio > zero)
                & (acceleration_ratio > zero)
                & (acceleration_ratio <= geodesic_acceptance_ratio)
            )
            loss_accelerated = jax.lax.cond(
                ratio_accepted,
                lambda _: jnp.sum(residual_value(theta + accelerated_step) ** 2),
                lambda _: jnp.asarray(jnp.inf, dtype=resid.dtype),
                operand=None,
            )
            used_geodesic = ratio_accepted & (loss_accelerated <= loss_velocity)
            step = jnp.where(used_geodesic, accelerated_step, velocity)
            step_sub = jnp.where(
                used_geodesic, velocity_sub + 0.5 * acceleration_sub, velocity_sub
            )
            loss_candidate = jnp.where(used_geodesic, loss_accelerated, loss_velocity)
        else:
            step, step_sub = velocity, velocity_sub
            loss_candidate = loss_velocity
            used_geodesic = jnp.asarray(False)
            acceleration_ratio = zero

        improved = jnp.isfinite(loss_candidate) & (loss_candidate < loss_old)
        theta_new = jnp.where(improved, theta + step, theta)
        damping_factor = jnp.where(improved, damping_decrease, damping_increase)
        new_damping = jnp.maximum(damping * damping_factor, damping_floor)
        loss = jnp.where(improved, loss_candidate, loss_old)

        # The input state's instances pass through verbatim -- None stays
        # None, so a user's own loop around update keeps its carry structure.
        instances = dict(metric=lm_state.metric, preconditioner=lm_state.preconditioner)
        if self.cache_jacobian:
            new_lm_state = LMState(
                new_damping,
                resid=resid,
                Jt=Jt,
                jacobian_valid=~improved,
                aux=aux,
                hyper=lm_state.hyper,
                solver_cache=step_solver.make_cache(~improved),
                **instances,
            )
        else:
            new_lm_state = LMState(new_damping, hyper=lm_state.hyper, **instances)
        return (
            unravel(theta_new),
            new_lm_state,
            LMInfo(
                loss=loss,
                loss_old=loss_old,
                loss_candidate=loss_candidate,
                accepted=improved,
                damping=new_damping,
                damping_factor=damping_factor,
                used_geodesic=used_geodesic,
                acceleration_ratio=acceleration_ratio,
                grad_norm=jnp.linalg.norm(step_solver.grad),
                step_norm=jnp.linalg.norm(step_sub),
                aux=aux,
            ),
        )

    def _converged(self, info, atol, gtol, xtol):
        atol_met = (atol > 0) & (jnp.sqrt(info.loss) < atol)
        gtol_met = (gtol > 0) & (info.grad_norm < gtol)
        xtol_met = (xtol > 0) & info.accepted & (info.step_norm < xtol)
        return atol_met | gtol_met | xtol_met

    def _cast_state(self, lm_state, dtype):
        return dataclasses.replace(
            lm_state,
            damping=jnp.asarray(lm_state.damping, dtype=dtype),
            hyper=_cast_hyper(lm_state.hyper, dtype),
        )

    def _ranking_objective(self, result, p, callback):
        # Without a callback info.loss already reports the objective at the
        # retained iterate; a callback can replace x/args after the last
        # update, so recompute. Nonfinite masks to +inf.
        if callback is None:
            loss = result.info.loss
        else:
            residual = self._residual_and_aux(result.x, result.args, p)[0]
            loss = jnp.sum(residual**2)
        return jnp.where(
            jnp.isfinite(loss), loss, jnp.asarray(jnp.inf, dtype=loss.dtype)
        )

    def _resolved_ad_solver(self, m, n):
        # The implicit-AD system is UNDAMPED, so the dual J~J~' is singular
        # whenever m > n and the normal J~'J~ whenever n > m. Each rule below
        # is only offered where its operator is invertible; SVD() covers the
        # cases where neither is (rank deficiency within the small side, as
        # padded zero residuals produce).
        resolved = self.ad_solver
        if resolved is None:
            # Match the forward family where its operator is invertible;
            # otherwise fall back to the assembled rule, which picks the
            # nonsingular side itself.
            if isinstance(self.linear_solver, GramCG) and m <= n:
                return self.linear_solver
            if isinstance(self.linear_solver, CG) and (
                n <= m or self.linear_solver.penalty is not None
            ):
                return self.linear_solver
            # Square: the tangent is a unique plain solve, so factor B itself
            # rather than squaring its condition number through B'B.
            if m == n:
                return LU()
            # Rectangular: the undamped system is singular whenever the small
            # side is rank deficient, which padded zero residuals produce by
            # construction, and the assembled Cholesky rules have no answer
            # there. SVD selects the minimum-metric-norm tangent at cond(B)
            # rather than cond(B)^2, so it is the safe default; Cholesky() is
            # the opt-in when the small side is known to have full rank.
            return SVD()
        if isinstance(resolved, LU) and m != n:
            raise ValueError(
                f"ad_solver=LU() needs a square system, but the residual is {m} "
                f"and x flattens to {n}: a rectangular tangent is selected by a "
                "minimum-metric-norm rule that a plain solve cannot express. "
                "Use SVD(), or Cholesky() when the small side has full rank"
            )
        if isinstance(resolved, GramCG) and m > n:
            raise ValueError(
                f"ad_solver=GramCG() needs m <= n, but the residual is {m} and "
                f"x flattens to {n}: the undamped dual J~J~' is singular there, "
                "so CG returns a wrong tangent rather than failing. Use SVD(), "
                "or CG(precond, penalty=...) to regularize"
            )
        if isinstance(resolved, CG) and n > m and resolved.penalty is None:
            raise ValueError(
                f"ad_solver=CG() needs n <= m, but the residual is {m} and x "
                f"flattens to {n}: the undamped normal J~'J~ is singular there, "
                "so CG returns a wrong tangent rather than failing. Use "
                "GramCG(precond), SVD(), or CG(precond, penalty=...)"
            )
        return resolved

    def _ad_x_tangent(self, x, args, p, p_dot, result, ad_success, initial_ad_point):
        if p is None:
            return jax.tree.map(_zero_tangent_leaf, x)
        # The carried instances are frozen conditioning data at the solution;
        # a failed lane reads the differentiation-inert pre-loop instances
        # instead (a callback may have left invalid arrays behind).
        lm_state = jax.lax.stop_gradient(result.lm_state)
        initial_instances = jax.lax.stop_gradient(initial_ad_point[3:5])
        lm_state = dataclasses.replace(
            lm_state,
            metric=_where_tree(ad_success, lm_state.metric, initial_instances[0]),
            preconditioner=_where_tree(
                ad_success, lm_state.preconditioner, initial_instances[1]
            ),
        )
        theta, unravel, residual, theta_jvp, residual_p_dot = self._ad_linearization(
            x, args, p, p_dot
        )
        resolved = self._resolved_ad_solver(residual.shape[0], theta.shape[0])
        if isinstance(resolved, LU):
            # Square: J theta_dot = -dr/dp p_dot has a unique solution, so the
            # metric selects nothing and the whitening round-trip below is
            # avoidable work -- solve the unwhitened system directly. This is
            # the same uniqueness that makes the plain factorization valid.
            Jt = self._assemble_jt(theta_jvp, theta, residual)
            return unravel(self._ad_tangent_lu(Jt, residual_p_dot))
        ctx = SolverContext(x=theta, lm_state=lm_state, args=args, p=p)
        n_m, n_f = self._block_sizes(theta.shape[0])
        dtype = residual.dtype

        def whiten(v):
            return jnp.asarray(self._extended_solve(v, ctx), dtype=dtype)

        def whiten_transpose(v):
            return jnp.asarray(self._extended_solve_transpose(v, ctx), dtype=dtype)

        if isinstance(resolved, (CG, GramCG)):
            u = self._ad_tangent_krylov(
                resolved,
                theta,
                theta_jvp,
                residual_p_dot,
                whiten,
                whiten_transpose,
                ctx,
            )
        else:
            # B' = F_bar^{-T} J', shape (n, m).
            Bt = whiten_transpose(self._assemble_jt(theta_jvp, theta, residual))
            u = (
                self._ad_tangent_svd(Bt, residual_p_dot)
                if isinstance(resolved, SVD)
                else self._ad_tangent_dense(Bt, residual_p_dot)
            )
        # S = F_bar^{-1} is not self-adjoint, and a matrix-free factor may
        # be opaque to JAX's transpose machinery, so declare its transpose
        # explicitly: the identity matvec exposes nothing to AD and every rule
        # routes through the declared solves. That keeps reverse mode working
        # through a metric JAX could not transpose on its own.
        theta_dot = jax.lax.custom_linear_solve(
            lambda v: v,
            u,
            lambda _, b: whiten(b),
            transpose_solve=lambda _, b: whiten_transpose(b),
        )
        return unravel(theta_dot)

    def _ad_tangent_dense(self, Bt, residual_p_dot):
        # Undamped Gauss-Newton tangent: u = -B^+ (dr/dp) p_dot through the
        # smaller of the two normal systems. Requires full rank; a rank-
        # deficient B needs SVD(), which selects the minimum-norm tangent.
        n, m = Bt.shape
        if n > m:
            factor = jsp_linalg.cho_factor(mm(Bt.T, Bt))
            return -mm(Bt, jsp_linalg.cho_solve(factor, residual_p_dot))
        factor = jsp_linalg.cho_factor(mm(Bt, Bt.T))
        return -jsp_linalg.cho_solve(factor, mm(Bt, residual_p_dot))

    def _ad_tangent_lu(self, Bt, residual_p_dot):
        # Square B: u = -B^{-1} (dr/dp) p_dot, factored directly. The normal
        # form (B'B)^{-1}B' is algebraically the same map here but numerically
        # worse, at cond(B)^2. Reverse mode transposes this solve, which is the
        # same factorization applied to B'.
        return -jnp.linalg.solve(Bt.T, residual_p_dot)

    def _ad_tangent_svd(self, Bt, residual_p_dot):
        # Spectral filter: u = -B^+ (dr/dp) p_dot, the minimum-metric-norm
        # tangent. This is the rule for the singular undamped systems that
        # padded zero residuals produce by construction, where the dense and
        # QR rules have no answer to give. The factors are constants in the
        # tangent program, so the map stays linear in residual_p_dot.
        U, sigma, Vt = jnp.linalg.svd(Bt.T, full_matrices=False)
        cutoff = max(Bt.shape) * jnp.finfo(Bt.dtype).eps * sigma[0]
        inverted = jnp.where(sigma > cutoff, 1.0 / jnp.maximum(sigma, cutoff), 0.0)
        return -mm(Vt.T, inverted * mm(U.T, residual_p_dot))

    def _ad_tangent_krylov(
        self, config, theta, theta_jvp, residual_p_dot, whiten, whiten_transpose, ctx
    ):
        dtype = residual_p_dot.dtype
        transpose_fn = jax.linear_transpose(theta_jvp, theta)
        zero_damping = jnp.zeros((), dtype=dtype)

        def JT(cotangent):
            return transpose_fn(cotangent)[0]

        def B(u):
            return theta_jvp(whiten(u))

        def Bt(w):
            return whiten_transpose(JT(w))

        apply_M = None
        ad_preconditioner = self._ad_preconditioner(ctx.lm_state)
        if ad_preconditioner is not None:
            # The AD system is undamped, so the preconditioner sees zero
            # damping (requires_positive_damping hooks were rejected).
            def apply_M(v):
                return ad_preconditioner.apply(v, zero_damping, ctx)

        def cg(matvec, rhs, preconditioner=apply_M):
            solution, _ = jsp_sparse_linalg.cg(
                matvec,
                rhs,
                tol=self._ad_cg_tol(dtype),
                atol=jnp.asarray(self.ad_solver_atol, dtype=dtype),
                maxiter=self.ad_solver_maxiter,
                M=preconditioner,
            )
            return solution

        if isinstance(config, GramCG):
            # Dual: (B B') y = (dr/dp) p_dot, then u = -B' y. Selection is
            # safe under any preconditioner here -- u = -B'y is invariant to
            # the null(B') component of y -- unlike the normal form, where the
            # preconditioner must preserve range(B').
            def dual_matvec(y):
                return B(Bt(y))

            y = jax.lax.custom_linear_solve(
                dual_matvec,
                residual_p_dot,
                lambda _, c: cg(dual_matvec, c),
                symmetric=True,
            )
            return -Bt(y)

        penalty = self.ad_solver_penalty

        def normal_matvec(u):
            value = Bt(B(u))
            if penalty is not None:
                value = value + jnp.asarray(penalty, dtype=dtype) * u
            return value

        if penalty is not None:
            # N = B'B + penalty I is SPD, so CG converges for any right-hand
            # side and the symmetric operator is its own transpose.
            transpose_solve = lambda _, c: cg(normal_matvec, c)  # noqa: E731
        else:
            # N = B'B is SINGULAR whenever B is (always, on the
            # underdetermined problems this solver targets). The forward
            # right-hand side lies in range(B') so CG converges there, but a
            # reverse-mode cotangent does not, and plain CG on it diverges.
            # Route the transpose through the push-through identity
            # N^+ = B' (BB')^{+2} B instead: each dual solve sees a right-hand
            # side in range(B), where BB' is invertible.
            def transpose_solve(_, c):
                return Bt(dual_solve(dual_solve(B(c))))

            def dual_solve(y):
                # UNPRECONDITIONED: this solve is posed on residual-space
                # m-vectors, while self.ad_solver_preconditioner is the
                # parameter-space one CG's own operator takes. Handing it an
                # m-vector is a shape error at best and a wrong tangent at
                # worst.
                return cg(lambda w: B(Bt(w)), y, preconditioner=None)

        rhs = -Bt(residual_p_dot)
        return jax.lax.custom_linear_solve(
            normal_matvec,
            rhs,
            lambda _, c: cg(normal_matvec, c),
            transpose_solve=transpose_solve,
        )

init(x0, args=None, *, p=None)

Build the initial :class:~nlls_gram.LMState at x0.

One residual evaluation types damping and sizes the Jacobian and linear-solver cache buffers. hyper stays None so manual update loops carry no extra buffers; solve populates it.

Source code in src/nlls_gram/metric_lm.py

def init(self, x0, args=None, *, p=None):
    """Build the initial :class:`~nlls_gram.LMState` at ``x0``.

    One residual evaluation types ``damping`` and sizes the Jacobian and
    linear-solver cache buffers. ``hyper`` stays ``None`` so manual
    ``update`` loops carry no extra buffers; ``solve`` populates it.
    """
    self._check_residual_args(args, p)
    residual, aux = self._residual_and_aux(x0, args, p)
    theta, _ = ravel_pytree(x0)
    n_m, _ = self._block_sizes(theta.size)
    dtype = residual.dtype
    damping = jnp.maximum(
        jnp.asarray(self.init_damping, dtype=dtype), _damping_floor(dtype)
    )
    instances = dict(
        metric=self.initial_metric, preconditioner=self.initial_preconditioner
    )
    if not self.cache_jacobian:
        return LMState(damping, **instances)
    return LMState(
        damping,
        resid=jnp.zeros(residual.shape, dtype=dtype),
        Jt=jnp.zeros((theta.size, residual.size), dtype=dtype),
        jacobian_valid=jnp.asarray(False, dtype=jnp.bool_),
        aux=jax.tree.map(jnp.zeros_like, aux),
        solver_cache=self.linear_solver.new_cache(
            residual.size, theta.size, n_m, dtype, False
        ),
        **instances,
    )

update(x, lm_state, args=None, p=None)

One LM step: returns (x_new, lm_state, info).

Source code in src/nlls_gram/metric_lm.py

def update(self, x, lm_state, args=None, p=None):
    """One LM step: returns ``(x_new, lm_state, info)``."""
    self._check_residual_args(args, p)
    theta, unravel = ravel_pytree(x)

    if self.has_aux:

        def residual_flat(th):
            value, aux = self.residual_fn(unravel(th), args, p)
            return jnp.ravel(value), aux

        def residual_value(th):
            return residual_flat(th)[0]

    else:

        def residual_flat(th):
            return jnp.ravel(self.residual_fn(unravel(th), args, p))

        residual_value = residual_flat

    jvp_fn = JT = Jt = None
    if not self.linear_solver.materializes_jacobian:
        if self.has_aux:
            resid, jvp_fn, aux = jax.linearize(residual_flat, theta, has_aux=True)
        else:
            resid, jvp_fn = jax.linearize(residual_flat, theta)
            aux = None
        transpose_fn = jax.linear_transpose(jvp_fn, theta)

        def JT(cotangent):
            return transpose_fn(cotangent)[0]

    elif self.cache_jacobian:
        resid, Jt, aux = jax.lax.cond(
            lm_state.jacobian_valid,
            lambda _: (lm_state.resid, lm_state.Jt, lm_state.aux),
            lambda _: self._dense_resid_jt_aux(residual_flat, theta),
            operand=None,
        )
    else:
        resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta)

    hyper = (
        lm_state.hyper
        if lm_state.hyper is not None
        else self.hyperparams(resid.dtype)
    )
    damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype)
    damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype)
    damping_floor = _damping_floor(resid.dtype)
    damping = jnp.maximum(
        jnp.asarray(lm_state.damping, dtype=resid.dtype), damping_floor
    )

    n_m, n_f = self._block_sizes(theta.shape[0])
    ctx = SolverContext(
        x=theta, lm_state=self._resolved_state(lm_state), args=args, p=p
    )
    zero = jnp.zeros((), dtype=resid.dtype)
    step_solver = self.linear_solver.prepare(
        Subproblem(
            resid=resid,
            theta=theta,
            Jt=Jt,
            jvp_fn=jvp_fn,
            JT=JT,
            whiten=lambda v: self._extended_solve(v, ctx),
            whiten_transpose=lambda v: self._extended_solve_transpose(v, ctx),
            y_m=jnp.zeros(n_m, dtype=resid.dtype),
            penalty_gradient=jnp.zeros(theta.shape[0], dtype=resid.dtype),
            ridge=zero,
            damping=damping,
            n_m=n_m,
            n_f=n_f,
            cache=lm_state.solver_cache,
            cache_enabled=self.cache_jacobian,
            hyper=hyper,
            ctx=ctx,
            penalized=False,
        )
    )

    # The solves produce the whitened step; the x-space step maps back
    # through the factor solve.
    velocity_sub = step_solver.velocity()
    velocity = jnp.asarray(self._extended_solve(velocity_sub, ctx), resid.dtype)
    loss_old = jnp.sum(resid**2)
    resid_velocity = residual_value(theta + velocity)
    loss_velocity = jnp.sum(resid_velocity**2)

    # Geodesic second-order correction, sharing the factorization.
    if self.geodesic_acceleration:
        geodesic_acceptance_ratio = jnp.asarray(
            hyper.geodesic_acceptance_ratio, dtype=resid.dtype
        )

        def first_jvp(th):
            # [1] is the tangent with and without has_aux.
            return jax.jvp(residual_flat, (th,), (velocity,), has_aux=self.has_aux)[
                1
            ]

        f_vv = jax.jvp(first_jvp, (theta,), (velocity,))[1]
        acceleration_sub = step_solver.correction(f_vv)
        acceleration = jnp.asarray(
            self._extended_solve(acceleration_sub, ctx), dtype=resid.dtype
        )
        accelerated_step = velocity + 0.5 * acceleration
        # The ratio criterion lives in the damping geometry's norm -- the
        # whitened one.
        acceleration_ratio = (
            2.0
            * jnp.linalg.norm(acceleration_sub)
            / (jnp.linalg.norm(velocity_sub) + jnp.finfo(resid.dtype).eps)
        )
        ratio_accepted = (
            (geodesic_acceptance_ratio > zero)
            & (acceleration_ratio > zero)
            & (acceleration_ratio <= geodesic_acceptance_ratio)
        )
        loss_accelerated = jax.lax.cond(
            ratio_accepted,
            lambda _: jnp.sum(residual_value(theta + accelerated_step) ** 2),
            lambda _: jnp.asarray(jnp.inf, dtype=resid.dtype),
            operand=None,
        )
        used_geodesic = ratio_accepted & (loss_accelerated <= loss_velocity)
        step = jnp.where(used_geodesic, accelerated_step, velocity)
        step_sub = jnp.where(
            used_geodesic, velocity_sub + 0.5 * acceleration_sub, velocity_sub
        )
        loss_candidate = jnp.where(used_geodesic, loss_accelerated, loss_velocity)
    else:
        step, step_sub = velocity, velocity_sub
        loss_candidate = loss_velocity
        used_geodesic = jnp.asarray(False)
        acceleration_ratio = zero

    improved = jnp.isfinite(loss_candidate) & (loss_candidate < loss_old)
    theta_new = jnp.where(improved, theta + step, theta)
    damping_factor = jnp.where(improved, damping_decrease, damping_increase)
    new_damping = jnp.maximum(damping * damping_factor, damping_floor)
    loss = jnp.where(improved, loss_candidate, loss_old)

    # The input state's instances pass through verbatim -- None stays
    # None, so a user's own loop around update keeps its carry structure.
    instances = dict(metric=lm_state.metric, preconditioner=lm_state.preconditioner)
    if self.cache_jacobian:
        new_lm_state = LMState(
            new_damping,
            resid=resid,
            Jt=Jt,
            jacobian_valid=~improved,
            aux=aux,
            hyper=lm_state.hyper,
            solver_cache=step_solver.make_cache(~improved),
            **instances,
        )
    else:
        new_lm_state = LMState(new_damping, hyper=lm_state.hyper, **instances)
    return (
        unravel(theta_new),
        new_lm_state,
        LMInfo(
            loss=loss,
            loss_old=loss_old,
            loss_candidate=loss_candidate,
            accepted=improved,
            damping=new_damping,
            damping_factor=damping_factor,
            used_geodesic=used_geodesic,
            acceleration_ratio=acceleration_ratio,
            grad_norm=jnp.linalg.norm(step_solver.grad),
            step_norm=jnp.linalg.norm(step_sub),
            aux=aux,
        ),
    )

nlls_gram.AnnealRidge dataclass

Ridge-continuation solve callback: multiply lm_state.ridge by decrease whenever the current level has yielded what it can, never below ridge_floor.

A level has yielded what it can when it is

• stationary: info.grad_norm fell below grad_rtol relative to its reference value at the current ridge level (the level's first observed gradient), or
• stalled (opt-in, stall_rtol > 0): an ACCEPTED step improved the gradient by less than a factor stall_rtol. The escape hatch for a frozen anneal: the per-level references COMPOUND -- with closely spaced levels each new reference is already small, the effective demand approaches grad_rtol ** levels times the initial gradient, and the anneal can freeze below the problem's noise floor while steps keep being accepted with negligible progress. Enable with stall_rtol ~ 0.99 when that happens. Off by default because it cannot distinguish a converged level from an accepted micro-step under temporarily high damping early in a hard solve -- where a false advance collapses the schedule prematurely; widening decrease (e.g. 0.01) is the alternative fix for a frozen anneal, since larger jumps keep the per-level references generous.

ridge_floor is REQUIRED and strictly positive (ridge = 0 is out of the solver's contract). Usage::

anneal = AnnealRidge(ridge_floor=1e-10)
result = solver.solve(x0, callback=anneal,
                      user_state=anneal.init_state(),
                      gtol=1e-8, atol=1e-8)

The solved-out continuation path converges to the minimum-seminorm solution by nonlinear Tikhonov theory (Engl-Kunisch-Neubauer 1989; Engl-Hanke-Neubauer 1996), while annealing per accepted stationarity event rather than per fully solved level is the iteratively regularized Gauss-Newton method (Bakushinskii 1992; Blaschke-Neubauer-Scherzer 1997; Kaltenbacher-Neubauer-Scherzer 2008), whose theory wants exactly this kind of monotone, boundedly geometric schedule. Pair the schedule with the conjunctive stopping rule: choose atol BETWEEN the ridge-floor residual and the last intermediate level's residual (they differ by roughly 1 / decrease), so the solve can only stop at the floor even when gtol must sit above a variant-dependent stationarity noise floor -- intermediate levels are stationary too, and atol is what rules them out.

HOW IT WORKS, for composing the schedule into a callback of your own (data re-draws, damping resets, a preconditioner refresh): the per-level reference and previous gradient ride in user_state as two fixed-shape scalars (init_state builds them; a lax.while_loop carry cannot grow from None mid-loop), with +inf marking "no observation at this level yet" -- the first step after a decrease sets the reference and can never read as stalled. Both trackers reset to +inf exactly when the level advances; at the floor the ridge stops changing, so the solver stops suppressing convergence and gtol/ atol can fire. All comparisons run at the ridge dtype and the returned trackers are cast back to the user_state dtype. A wrapping callback calls the instance, inspects action.lm_state.ridge < ctx.lm_state.ridge for a level advance, and returns a dataclasses.replace of the action -- e.g. gating an expensive preconditioner rebuild on the advance with jax.lax.cond so the eigendecomposition is paid only when it fires (a jnp.where merge would pay it every step)::

def driver_callback(ctx):
    action = anneal(ctx)
    advanced = action.lm_state.ridge < ctx.lm_state.ridge
    precond = jax.lax.cond(
        advanced,
        lambda: BlockEigenPreconditioner(families(ctx.x), PERM),
        lambda: ctx.lm_state.preconditioner,
    )
    return dataclasses.replace(
        action,
        lm_state=dataclasses.replace(
            action.lm_state, preconditioner=precond
        ),
    )

A frozen dataclass with scalar fields rather than a closure because solve marks the callback a jit STATIC argument: equal schedules compare equal and share one compiled loop, while a fresh closure per construction would key a fresh compilation. ridge_floor must be a concrete float for that sharing -- a traced value is unhashable and falls back to identity.

Source code in src/nlls_gram/ridge_lm.py

@dataclass(frozen=True)
class AnnealRidge:
    """Ridge-continuation ``solve`` callback: multiply ``lm_state.ridge`` by
    ``decrease`` whenever the current level has yielded what it can, never
    below ``ridge_floor``.

    A level has yielded what it can when it is

    - **stationary**: ``info.grad_norm`` fell below ``grad_rtol`` relative to
      its reference value at the current ridge level (the level's first
      observed gradient), or
    - **stalled** (opt-in, ``stall_rtol > 0``): an ACCEPTED step improved
      the gradient by less than a factor ``stall_rtol``. The escape hatch
      for a frozen anneal: the per-level references COMPOUND -- with closely
      spaced levels each new reference is already small, the effective
      demand approaches ``grad_rtol ** levels`` times the initial gradient,
      and the anneal can freeze below the problem's noise floor while steps
      keep being accepted with negligible progress. Enable with
      ``stall_rtol ~ 0.99`` when that happens. Off by default because it
      cannot distinguish a converged level from an accepted micro-step under
      temporarily high damping early in a hard solve -- where a false
      advance collapses the schedule prematurely; widening ``decrease``
      (e.g. ``0.01``) is the alternative fix for a frozen anneal, since
      larger jumps keep the per-level references generous.

    ``ridge_floor`` is REQUIRED and strictly positive (``ridge = 0`` is out
    of the solver's contract). Usage::

        anneal = AnnealRidge(ridge_floor=1e-10)
        result = solver.solve(x0, callback=anneal,
                              user_state=anneal.init_state(),
                              gtol=1e-8, atol=1e-8)

    The solved-out continuation path converges to the minimum-seminorm
    solution by nonlinear Tikhonov theory (Engl-Kunisch-Neubauer 1989;
    Engl-Hanke-Neubauer 1996), while annealing per accepted stationarity
    event rather than per fully solved level is the iteratively regularized
    Gauss-Newton method (Bakushinskii 1992; Blaschke-Neubauer-Scherzer 1997;
    Kaltenbacher-Neubauer-Scherzer 2008), whose theory wants exactly this
    kind of monotone, boundedly geometric schedule. Pair the schedule with
    the conjunctive stopping rule: choose ``atol`` BETWEEN the ridge-floor
    residual and the last intermediate level's residual (they differ by
    roughly ``1 / decrease``), so the solve can only stop at the floor even
    when ``gtol`` must sit above a variant-dependent stationarity noise
    floor -- intermediate levels are stationary too, and ``atol`` is what
    rules them out.

    HOW IT WORKS, for composing the schedule into a callback of your own
    (data re-draws, damping resets, a preconditioner refresh): the per-level
    reference and previous gradient ride in ``user_state`` as two
    fixed-shape scalars (``init_state`` builds them; a ``lax.while_loop``
    carry cannot grow from ``None`` mid-loop), with ``+inf`` marking "no
    observation at this level yet" -- the first step after a decrease sets
    the reference and can never read as stalled. Both trackers reset to
    ``+inf`` exactly when the level advances; at the floor the ridge stops
    changing, so the solver stops suppressing convergence and ``gtol``/
    ``atol`` can fire. All comparisons run at the ridge dtype and the
    returned trackers are cast back to the ``user_state`` dtype. A wrapping
    callback calls the instance, inspects ``action.lm_state.ridge <
    ctx.lm_state.ridge`` for a level advance, and returns a
    ``dataclasses.replace`` of the action -- e.g. gating an expensive
    preconditioner rebuild on the advance with ``jax.lax.cond`` so the
    eigendecomposition is paid only when it fires (a ``jnp.where`` merge
    would pay it every step)::

        def driver_callback(ctx):
            action = anneal(ctx)
            advanced = action.lm_state.ridge < ctx.lm_state.ridge
            precond = jax.lax.cond(
                advanced,
                lambda: BlockEigenPreconditioner(families(ctx.x), PERM),
                lambda: ctx.lm_state.preconditioner,
            )
            return dataclasses.replace(
                action,
                lm_state=dataclasses.replace(
                    action.lm_state, preconditioner=precond
                ),
            )

    A frozen dataclass with scalar fields rather than a closure because
    ``solve`` marks the callback a jit STATIC argument: equal schedules
    compare equal and share one compiled loop, while a fresh closure per
    construction would key a fresh compilation. ``ridge_floor`` must be a
    concrete float for that sharing -- a traced value is unhashable and
    falls back to identity.
    """

    ridge_floor: float
    decrease: float = 0.1
    grad_rtol: float = 1e-2
    stall_rtol: float = 0.0

    def __post_init__(self):
        if not 0 < self.decrease < 1:
            raise ValueError("decrease must lie strictly between 0 and 1")
        if not isinstance(self.ridge_floor, (jax.Array, jax.core.Tracer)) and (
            float(self.ridge_floor) <= 0
        ):
            raise ValueError(
                "ridge_floor must be strictly positive (ridge = 0 is "
                "unsupported by RidgeLevenbergMarquardt)"
            )
        if self.grad_rtol <= 0:
            raise ValueError("grad_rtol must be positive")
        if not 0 <= self.stall_rtol < 1:
            raise ValueError("stall_rtol must lie in [0, 1)")

    def init_state(self, dtype=None):
        """The initial ``user_state``: both trackers at ``+inf``. ``dtype``
        defaults to the JAX default float; pass the problem dtype explicitly
        for a float32 program under enabled x64."""
        infinity = jnp.asarray(
            jnp.inf, dtype=jnp.result_type(float) if dtype is None else dtype
        )
        return {"reference": infinity, "previous": infinity}

    def __call__(self, ctx):
        ridge = ctx.lm_state.ridge
        dtype = ridge.dtype
        grad_norm = jnp.asarray(ctx.info.grad_norm, dtype=dtype)
        reference = jnp.asarray(ctx.user_state["reference"], dtype=dtype)
        previous = jnp.asarray(ctx.user_state["previous"], dtype=dtype)
        # +inf marks "no observation at this level yet": the first step after a
        # ridge decrease (or the initial step) sets the reference and can never
        # read as stalled.
        reference = jnp.where(jnp.isfinite(reference), reference, grad_norm)
        stationary = grad_norm <= jnp.asarray(self.grad_rtol, dtype) * reference
        if self.stall_rtol > 0:
            # ACCEPTED steps only: a rejected step leaves x (and so the
            # gradient) unchanged -- that is the trust region adapting, not the
            # level converging -- while an accepted step that improved the
            # gradient by less than the stall factor means the level has
            # yielded what it can.
            stalled = ctx.info.accepted & (
                grad_norm >= jnp.asarray(self.stall_rtol, dtype) * previous
            )
        else:
            stalled = jnp.asarray(False)
        new_ridge = jnp.where(
            stationary | stalled,
            jnp.maximum(
                ridge * jnp.asarray(self.decrease, dtype),
                jnp.asarray(self.ridge_floor, dtype),
            ),
            ridge,
        )
        # Reset the trackers when the level actually changes; at the floor the
        # ridge is unchanged, so convergence is not suppressed and gtol/atol
        # can fire.
        advanced = new_ridge < ridge
        fresh_level = jnp.asarray(jnp.inf, dtype)
        state_dtype = jnp.asarray(ctx.user_state["reference"]).dtype
        return LMAction(
            lm_state=dataclasses.replace(ctx.lm_state, ridge=new_ridge),
            user_state={
                "reference": jnp.where(advanced, fresh_level, reference).astype(
                    state_dtype
                ),
                "previous": jnp.where(advanced, fresh_level, grad_norm).astype(
                    state_dtype
                ),
            },
        )

init_state(dtype=None)

The initial user_state: both trackers at +inf. dtype defaults to the JAX default float; pass the problem dtype explicitly for a float32 program under enabled x64.

Source code in src/nlls_gram/ridge_lm.py

def init_state(self, dtype=None):
    """The initial ``user_state``: both trackers at ``+inf``. ``dtype``
    defaults to the JAX default float; pass the problem dtype explicitly
    for a float32 program under enabled x64."""
    infinity = jnp.asarray(
        jnp.inf, dtype=jnp.result_type(float) if dtype is None else dtype
    )
    return {"reference": infinity, "previous": infinity}

Linear solvers

nlls_gram.Cholesky dataclass

Bases: LinearSolver

Dense factorization of the damped subproblem.

form="normal" factors the n x n whitened normal system G = J~'J~ + ridge E (cached across rejected steps, where only the damping changed, so a reject pays the n^3/3 refactor without the GEMM and without re-materializing J~') and solves (G + damping I) u = -g.

form="gram" factors the m x m dual D = J~ J~' instead and takes the step u = -J~'(D + damping I)^{-1} r. For damping > 0 the two produce the same step in exact arithmetic, by the push-through identity B'(BB' + lam I)^{-1} = (B'B + lam I)^{-1}B'.

In FLOATING POINT they part company when m > n: the dual then carries m - n structural zero eigenvalues, so its condition number is sigma_max^2 / damping however well conditioned J~ is, and the gram step loses digits as damping falls (measured ~1e-2 relative at damping=1e-14 where the normal form holds 1e-15). form="auto" (the default) never picks gram there -- it takes gram only when n > m, where the normal system is the singular one -- so the default is safe and the choice is a cost one. Forcing form="gram" on a tall problem is not.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class Cholesky(LinearSolver):
    """Dense factorization of the damped subproblem.

    ``form="normal"`` factors the ``n x n`` whitened normal system
    ``G = J~'J~ + ridge E`` (cached across rejected steps, where only the
    damping changed, so a reject pays the ``n^3/3`` refactor without the GEMM
    and without re-materializing ``J~'``) and solves
    ``(G + damping I) u = -g``.

    ``form="gram"`` factors the ``m x m`` dual ``D = J~ J~'`` instead and
    takes the step ``u = -J~'(D + damping I)^{-1} r``. For ``damping > 0`` the
    two produce the same step in exact arithmetic, by the push-through
    identity ``B'(BB' + lam I)^{-1} = (B'B + lam I)^{-1}B'``.

    In FLOATING POINT they part company when ``m > n``: the dual then carries
    ``m - n`` structural zero eigenvalues, so its condition number is
    ``sigma_max^2 / damping`` however well conditioned ``J~`` is, and the gram
    step loses digits as damping falls (measured ~1e-2 relative at
    ``damping=1e-14`` where the normal form holds 1e-15). ``form="auto"`` (the
    default) never picks gram there -- it takes gram only when ``n > m``, where
    the normal system is the singular one -- so the default is safe and the
    choice is a cost one. Forcing ``form="gram"`` on a tall problem is not.
    """

    form: str = "auto"

    def __post_init__(self):
        if self.form not in ("auto", "gram", "normal"):
            raise ValueError(
                f"Cholesky.form must be 'auto', 'gram', or 'normal'; got {self.form!r}"
            )

    @property
    def supports_penalty(self):
        # form="auto" resolves to normal under a penalized subproblem.
        return self.form != "gram"

    def _resolved_form(self, m, n, penalized):
        # The ridge solver's penalty rows have no dual analogue -- the dual
        # operator J~J~' never sees them -- so a penalized subproblem is
        # always the normal form. The ridge constructor rejects an explicit
        # form="gram" rather than silently ignoring it.
        if penalized:
            return "normal"
        if self.form != "auto":
            return self.form
        return "gram" if n > m else "normal"

    def new_cache(self, m, n, n_m, dtype, penalized):
        size = m if self._resolved_form(m, n, penalized) == "gram" else n
        return CholeskyCache(
            G=jnp.zeros((size, size), dtype=dtype),
            valid=jnp.asarray(False, dtype=jnp.bool_),
            ridge=jnp.zeros((), dtype=dtype),
        )

    def prepare(self, sub):
        n_m, ridge = sub.n_m, sub.ridge
        # B' = F_bar^{-T} J', shape (n, m). Every form below is built from it.
        grad = sub.whitened_transpose(mm(sub.Jt, sub.resid))
        if sub.penalized:
            grad = grad + ridge * sub.penalty_gradient
        gram = self._resolved_form(sub.m, sub.n, sub.penalized) == "gram"

        def assemble():
            Bt = sub.whitened_transpose(sub.Jt)
            if gram:
                return mm(Bt.T, Bt)
            normal = mm(Bt, Bt.T)
            if not sub.penalized:
                return normal
            diagonal = jnp.arange(n_m)
            return normal.at[diagonal, diagonal].add(ridge)

        matrix = sub.cached(assemble)
        size = sub.m if gram else sub.n
        shift = jnp.arange(size)
        factor = jsp_linalg.cho_factor(matrix.at[shift, shift].add(sub.damping))
        if gram:
            # u = -B'(D + damping I)^{-1} c on residual-space right-hand sides.
            def dual_step(c):
                return -sub.whitened_transpose(
                    mm(sub.Jt, jsp_linalg.cho_solve(factor, c))
                )

            velocity, correction = (lambda: dual_step(sub.resid)), dual_step
        else:

            def normal_step(c):
                return -jsp_linalg.cho_solve(factor, c)

            velocity = lambda: normal_step(grad)  # noqa: E731
            correction = lambda f_vv: normal_step(  # noqa: E731
                sub.whitened_transpose(mm(sub.Jt, f_vv))
            )
        return StepSolver(
            grad=grad,
            velocity=velocity,
            correction=correction,
            make_cache=lambda valid: CholeskyCache(matrix, valid, ridge),
        )

nlls_gram.QR dataclass

Bases: LinearSolver

Damping-row QR of the augmented whitened stack.

One QR of [J~; sqrt(ridge) [I 0] | b~] (the penalty rows and the b~ tail only for the ridge solver) is cached per (x, ridge): its leading columns are the stack's R factor and its last column carries Q'b, so the velocity is a backward-stable least-squares solve with NO normal equations -- More 1978's damping-row structure, accurate at cond(A) rather than cond(A)^2. Each step re-factors only the damping rows, and those rows keep the system full rank for any damping > 0, so a rank-deficient Jacobian is handled rather than producing a non-finite step. No knobs.

Forward only: the implicit-AD system is undamped, so the damping rows that make this path well posed vanish there. Use SVD() for a rank-deficient tangent.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class QR(LinearSolver):
    """Damping-row QR of the augmented whitened stack.

    One QR of ``[J~; sqrt(ridge) [I 0] | b~]`` (the penalty rows and the
    ``b~`` tail only for the ridge solver) is cached per ``(x, ridge)``: its
    leading columns are the stack's R factor and its last column carries
    ``Q'b``, so the velocity is a backward-stable least-squares solve with NO
    normal equations -- More 1978's damping-row structure, accurate at
    ``cond(A)`` rather than ``cond(A)^2``. Each step re-factors only the
    damping rows, and those rows keep the system full rank for any
    ``damping > 0``, so a rank-deficient Jacobian is handled rather than
    producing a non-finite step. No knobs.

    Forward only: the implicit-AD system is undamped, so the damping rows that
    make this path well posed vanish there. Use ``SVD()`` for a rank-deficient
    tangent.
    """

    supports_ad = False

    def new_cache(self, m, n, n_m, dtype, penalized):
        rows = min(m + n_m, n + 1) if penalized else min(m, n + 1)
        return QRCache(
            R=jnp.zeros((rows, n + 1), dtype=dtype),
            valid=jnp.asarray(False, dtype=jnp.bool_),
            ridge=jnp.zeros((), dtype=dtype),
        )

    def prepare(self, sub):
        n, n_m, ridge, dtype = sub.n, sub.n_m, sub.ridge, sub.dtype
        sqrt_ridge = jnp.sqrt(ridge)
        grad = sub.whitened_transpose(mm(sub.Jt, sub.resid))
        if sub.penalized:
            grad = grad + ridge * sub.penalty_gradient

        def assemble():
            Bt = sub.whitened_transpose(sub.Jt)
            rows, rhs = [Bt.T], [sub.resid]
            if sub.penalized:
                rows.append(sqrt_ridge * jnp.eye(n_m, n, dtype=dtype))
                rhs.append(sqrt_ridge * sub.y_m)
            stacked = jnp.concatenate(rows, axis=0)
            b_stacked = jnp.concatenate(rhs)
            return jnp.linalg.qr(
                jnp.concatenate([stacked, b_stacked[:, None]], axis=1), mode="r"
            )

        qr_R = sub.cached(assemble)
        r_factor, transformed_rhs = qr_R[:, :-1], qr_R[:, -1]
        # Per-step damping-row refactor: [R; sqrt(damping) I] = Q2 R2 with
        # R2'R2 = A'A + damping I. When the cached R is upper trapezoidal these
        # rows are what make the system full rank. Q2 is retained to transform
        # the velocity right-hand side stably.
        damped_stack = jnp.concatenate(
            [r_factor, jnp.sqrt(sub.damping) * jnp.eye(n, dtype=dtype)], axis=0
        )
        Q_mu, R_mu = jnp.linalg.qr(damped_stack, mode="reduced")

        def damped_normal_matvec(v):
            gauss_newton = sub.whitened_transpose(
                mm(sub.Jt, mm(sub.Jt.T, sub.whitened(v)))
            )
            shift = sub.damping * v
            if sub.penalized:
                shift = shift + ridge * jnp.concatenate(
                    [v[:n_m], jnp.zeros(sub.n_f, dtype=dtype)]
                )
            return gauss_newton + shift

        def correction(f_vv):
            # Corrected semi-normal equations (Bjorck 1987): triangular solves
            # against R_mu, then ONE fixed iterative-refinement pass through
            # matvecs (Bjorck 1996 Sec. 6.6.5). The second-order correction
            # tolerates the squared conditioning; accept/reject guards it.
            b = -sub.whitened_transpose(mm(sub.Jt, f_vv))
            half = jsp_linalg.solve_triangular(R_mu.T, b, lower=True)
            delta = jsp_linalg.solve_triangular(R_mu, half, lower=False)
            correction_rhs = b - damped_normal_matvec(delta)
            half = jsp_linalg.solve_triangular(R_mu.T, correction_rhs, lower=True)
            return delta + jsp_linalg.solve_triangular(R_mu, half, lower=False)

        def velocity():
            # min ||[R; sqrt(damping) I] delta + [Q'b; 0]||^2 solved through
            # Q2: exact and backward stable at cond(A), never cond(A)^2.
            rhs = jnp.concatenate([transformed_rhs, jnp.zeros(n, dtype=dtype)])
            return -jsp_linalg.solve_triangular(R_mu, mm(Q_mu.T, rhs), lower=False)

        return StepSolver(
            grad=grad,
            velocity=velocity,
            correction=correction,
            make_cache=lambda valid: QRCache(qr_R, valid, ridge),
        )

nlls_gram.CG dataclass

Bases: _KrylovConfig

Matrix-free preconditioned CG on the whitened NORMAL operator, in parameter space.

As linear_solver it solves the damped forward subproblem (J~'J~ + ridge E + damping I) delta_y = -g -- the same SPD system :class:Cholesky factors -- with the preconditioner in CG's M slot at the live damping. As ad_solver it solves the undamped implicit-AD system, with the preconditioner applied at zero damping (subclasses marked requires_positive_damping are rejected for that role) and penalty optionally adding a small ridge that stabilizes a rank-deficient tangent.

preconditioner is REQUIRED in the forward role -- nobody should run Krylov methods without a preconditioning decision, so :class:~nlls_gram.IdentityPreconditioner is the explicit opt-out and a custom one is a small registered dataclass implementing apply(v, damping, ctx). In the ad_solver role, preconditioner=None inherits the CARRIED forward instance at the solution while pinning the AD tolerance and budget. On rank-deficient problems the preconditioner must map range(B') into itself or the minimum-norm selection is silently lost; the identity, polynomials in the operator, and exact shifted inverses are safe, and on full-column-rank problems the condition is vacuous.

tol=None resolves to a dtype default (1e-10 in float64, 1e-6 in float32); maxiter must be set when both tolerances are explicitly zero, since an uncapped zero-tolerance CG loop has no stopping rule.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class CG(_KrylovConfig):
    """Matrix-free preconditioned CG on the whitened NORMAL operator, in
    parameter space.

    As ``linear_solver`` it solves the damped forward subproblem
    ``(J~'J~ + ridge E + damping I) delta_y = -g`` -- the same SPD system
    :class:`Cholesky` factors -- with the ``preconditioner`` in CG's ``M`` slot
    at the live damping. As ``ad_solver`` it solves the undamped implicit-AD
    system, with the preconditioner applied at zero damping (subclasses marked
    ``requires_positive_damping`` are rejected for that role) and ``penalty``
    optionally adding a small ridge that stabilizes a rank-deficient tangent.

    ``preconditioner`` is REQUIRED in the forward role -- nobody should run
    Krylov methods without a preconditioning decision, so
    :class:`~nlls_gram.IdentityPreconditioner` is the explicit opt-out and a
    custom one is a small registered dataclass implementing
    ``apply(v, damping, ctx)``. In the ``ad_solver`` role,
    ``preconditioner=None`` inherits the CARRIED forward instance at the
    solution while pinning the AD tolerance and budget. On rank-deficient
    problems the preconditioner must map ``range(B')`` into itself or the
    minimum-norm selection is silently lost; the identity, polynomials in the
    operator, and exact shifted inverses are safe, and on full-column-rank
    problems the condition is vacuous.

    ``tol=None`` resolves to a dtype default (``1e-10`` in float64, ``1e-6``
    in float32); ``maxiter`` must be set when both tolerances are explicitly
    zero, since an uncapped zero-tolerance CG loop has no stopping rule.
    """

    preconditioner: Preconditioner | None
    tol: float | None = None
    atol: float = 0.0
    maxiter: int | None = None
    penalty: float | None = None

    def prepare(self, sub):
        n_m, n_f, ridge, dtype = sub.n_m, sub.n_f, sub.ridge, sub.dtype
        damping, ctx = sub.damping, sub.ctx
        sqrt_ridge = jnp.sqrt(ridge)

        # Whitened operator J~ = J F_bar^{-1}: products route through the
        # metric's factor callbacks.
        def J_sub(u):
            return sub.jvp_fn(sub.whitened(u))

        def JT_sub(w):
            return sub.whitened_transpose(sub.JT(w))

        grad = JT_sub(sub.resid)
        if sub.penalized:
            grad = grad + ridge * sub.penalty_gradient

        # N = A'A + damping I for the augmented A = [J~; sqrt(ridge) [I 0]] --
        # the preconditioner-free SPD operator that custom_linear_solve
        # differentiates through, posed on u.
        def N_matvec(u):
            normal = JT_sub(J_sub(u))
            if sub.penalized:
                pullback = sqrt_ridge * jnp.concatenate(
                    [sqrt_ridge * u[:n_m], jnp.zeros(n_f, dtype=dtype)]
                )
                normal = normal + pullback
            return normal + damping * u

        # The CARRIED instance, so a callback refresh reaches the very next
        # inner solve; the config's own field only seeds the initial state.
        precond = ctx.lm_state.preconditioner

        def apply_M(v):
            return precond.apply(v, damping, ctx)

        if precond is None:
            apply_M = None

        def solve_N(_, c):
            return self._cg(N_matvec, c, sub, apply_M)

        def solve(c):
            return jax.lax.custom_linear_solve(
                N_matvec, -c, solve=solve_N, transpose_solve=solve_N, symmetric=True
            )

        return StepSolver(
            grad=grad,
            velocity=lambda: solve(grad),
            correction=lambda f_vv: solve(JT_sub(f_vv)),
            make_cache=lambda valid: None,
        )

nlls_gram.GramCG dataclass

Bases: _KrylovConfig

Matrix-free preconditioned CG on the DUAL operator, in residual space.

Applies y -> J~ J~' y + damping y on m-vectors and takes the step u = -J~'y, so the Krylov iteration lives in residual dimension -- the matrix-free form for the m << n regime this package targets, where :class:CG's n-dimensional iteration is the expensive one. At inner convergence the step matches :class:CG's, and a budget-truncated step still lies in range(J~'), so the minimum-metric-norm structure survives truncation.

preconditioner acts on residual-space vectors -- an SPD approximation of (J~J~' + damping I)^{-1} -- which is the only difference from :class:CG's contract, and the reason the two are separate configs rather than one with a flag. The ridge penalty rows have no dual analogue, so this config serves :class:~nlls_gram.LevenbergMarquardt alone.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class GramCG(_KrylovConfig):
    """Matrix-free preconditioned CG on the DUAL operator, in residual space.

    Applies ``y -> J~ J~' y + damping y`` on ``m``-vectors and takes the step
    ``u = -J~'y``, so the Krylov iteration lives in residual dimension -- the
    matrix-free form for the ``m << n`` regime this package targets, where
    :class:`CG`'s ``n``-dimensional iteration is the expensive one. At inner
    convergence the step matches :class:`CG`'s, and a budget-truncated step
    still lies in ``range(J~')``, so the minimum-metric-norm structure
    survives truncation.

    ``preconditioner`` acts on residual-space vectors -- an SPD approximation
    of ``(J~J~' + damping I)^{-1}`` -- which is the only difference from
    :class:`CG`'s contract, and the reason the two are separate configs rather
    than one with a flag. The ridge penalty rows have no dual analogue, so
    this config serves :class:`~nlls_gram.LevenbergMarquardt` alone.
    """

    supports_penalty = False

    preconditioner: Preconditioner | None
    tol: float | None = None
    atol: float = 0.0
    maxiter: int | None = None

    def prepare(self, sub):
        damping, ctx = sub.damping, sub.ctx

        def dual_matvec(y):
            return sub.jvp_fn(sub.whitened(sub.whitened_transpose(sub.JT(y)))) + (
                damping * y
            )

        # The CARRIED instance, so a callback refresh reaches the very next
        # inner solve; the config's own field only seeds the initial state.
        precond = ctx.lm_state.preconditioner

        def apply_M(v):
            return precond.apply(v, damping, ctx)

        if precond is None:
            apply_M = None

        def solve_dual(_, c):
            return self._cg(dual_matvec, c, sub, apply_M)

        def step(c):
            y = jax.lax.custom_linear_solve(
                dual_matvec,
                c,
                solve=solve_dual,
                transpose_solve=solve_dual,
                symmetric=True,
            )
            return -sub.whitened_transpose(sub.JT(y))

        return StepSolver(
            grad=sub.whitened_transpose(sub.JT(sub.resid)),
            velocity=lambda: step(sub.resid),
            correction=step,
            make_cache=lambda valid: None,
        )

nlls_gram.LU dataclass

Bases: LinearSolver

Direct nonsymmetric solve of a SQUARE system, ad_solver role only.

When the whitened Jacobian is square the implicit-AD system B u = -dr/dp has a unique solution, so there is no norm being minimized and the metric does not select among answers. That makes the plain factorization of B available, and it is strictly better than routing through a normal or dual operator: cond(B'B) = cond(B)^2, so Cholesky loses half the significant digits of the tangent for nothing. One factorization also serves both directions -- forward mode solves with B, reverse mode with B'.

This is why ad_solver=None resolves here for a square dense problem. It requires squareness and says so: a rectangular system needs a selection rule (minimum-metric-norm) that a plain solve cannot express, so :class:SVD or the Cholesky forms own that case.

The forward subproblem is damped and therefore symmetric positive definite at every shape, so this config has no forward role.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class LU(LinearSolver):
    """Direct nonsymmetric solve of a SQUARE system, ``ad_solver`` role only.

    When the whitened Jacobian is square the implicit-AD system ``B u = -dr/dp``
    has a unique solution, so there is no norm being minimized and the metric
    does not select among answers. That makes the plain factorization of ``B``
    available, and it is strictly better than routing through a normal or dual
    operator: ``cond(B'B) = cond(B)^2``, so Cholesky loses half the significant
    digits of the tangent for nothing. One factorization also serves both
    directions -- forward mode solves with ``B``, reverse mode with ``B'``.

    This is why ``ad_solver=None`` resolves here for a square dense problem.
    It requires squareness and says so: a rectangular system needs a selection
    rule (minimum-metric-norm) that a plain solve cannot express, so
    :class:`SVD` or the Cholesky forms own that case.

    The forward subproblem is damped and therefore symmetric positive definite
    at every shape, so this config has no forward role.
    """

    supports_forward = False
    supports_penalty = False

    def prepare(self, sub):
        raise NotImplementedError("LU is an ad_solver, not a forward solver")

nlls_gram.SVD dataclass

Bases: LinearSolver

Spectral-filter pseudoinverse, for the ad_solver role only.

The implicit-AD system is UNDAMPED, so it is singular whenever the whitened Jacobian is rank deficient -- padded zero residuals make it so by construction. This rule truncates at max(m, n) * eps * sigma_max and returns the minimum-metric-norm tangent, which is the right answer there rather than a NaN or a silent pseudo-solve. It assembles, so it is the dense fallback rather than the default.

Source code in src/nlls_gram/linear_solvers.py

@dataclass(frozen=True)
class SVD(LinearSolver):
    """Spectral-filter pseudoinverse, for the ``ad_solver`` role only.

    The implicit-AD system is UNDAMPED, so it is singular whenever the
    whitened Jacobian is rank deficient -- padded zero residuals make it so by
    construction. This rule truncates at ``max(m, n) * eps * sigma_max`` and
    returns the minimum-metric-norm tangent, which is the right answer there
    rather than a NaN or a silent pseudo-solve. It assembles, so it is the
    dense fallback rather than the default.
    """

    supports_forward = False
    supports_penalty = False

    def prepare(self, sub):
        raise NotImplementedError("SVD is an ad_solver, not a forward solver")

Metrics

nlls_gram.Metric

Positive-definite metric W, via factor callbacks.

x = [x_m; x_f] splits into the metric block x_m -- the leading size coordinates, covered by W -- and a free block x_f. The metric is supplied through callbacks for an invertible factor F with W = F'F, the whitened variable being F x_m and ||v||_W = ||F v||_2. F should be upper triangular; the canonical example is F = jnp.linalg.cholesky(K, upper=True). The solver extends it to F_bar = blockdiag(F, sqrt(free_scale) I) over the whole vector and never materializes either.

Subclasses implement the ops on metric-block vectors (or matrices whose LEADING axis is size; columns are batched):

• factor_apply(v, ctx): F v
• factor_solve(v, ctx): F^{-1} v
• factor_solve_transpose(v, ctx): F^{-T} v
• norm(v, ctx): ||v||_W, defaulted to ||factor_apply(v)||_2. Provided for callers; the solvers measure in the whitened variable and never call it.

Every op receives a :class:~nlls_gram.SolverContext carrying the solver's live state, so an exotic metric can key off the iterate through ctx.x and ctx.lm_state.

Metric instances are JAX PYTREES: array fields are traced leaves, the type plus its static fields are structure. Every concrete class must be registered with :func:~nlls_gram.register_pytree_dataclass -- the solvers reject unregistered instances. The instance rides inside the solver state (lm_state.metric), so a solve callback replaces the metric by constructing a new instance of the same type (same static fields, same leaf shapes and dtypes) -- pure traced ops, no recompilation. Equal-config instances with fresh arrays share one compiled solve loop. dataclasses.replace works too: metric constructors only validate and derive shapes, so re-running them under trace is cheap.

free_scale weights the free block in the whitened variable: 1.0 (the default) leaves it Euclidean. The ridge solver never penalizes the free block whatever the scale -- free_scale only changes its trust-region geometry -- while for the metric solver it IS that block's damping weight. It is a traced leaf, canonicalized by each constructor to the factor's float dtype, so changing it never recompiles.

Contracts: the factor must be EXACT. The solver hardcodes the identity penalty block in the whitened variable, so an approximate factor silently changes the objective (unlike a CG preconditioner, which may be sloppy). The ridge weight never enters the factorization, so ridge continuation composes unchanged. How a subclass fulfills the ops -- prefactorized storage, factorize-in-__init__, fully matrix-free -- is its constructor's business, and the constructor must be traceable when the metric is rebuilt inside a jitted callback.

Source code in src/nlls_gram/metrics.py

class Metric:
    """Positive-definite metric ``W``, via factor callbacks.

    ``x = [x_m; x_f]`` splits into the metric block ``x_m`` -- the leading
    ``size`` coordinates, covered by ``W`` -- and a free block ``x_f``. The
    metric is supplied through callbacks for an invertible factor ``F`` with
    ``W = F'F``, the whitened variable being ``F x_m`` and
    ``||v||_W = ||F v||_2``. ``F`` should be upper triangular; the canonical
    example is ``F = jnp.linalg.cholesky(K, upper=True)``. The solver extends
    it to ``F_bar = blockdiag(F, sqrt(free_scale) I)`` over the whole vector
    and never materializes either.

    Subclasses implement the ops on metric-block vectors (or matrices whose
    LEADING axis is ``size``; columns are batched):

    - ``factor_apply(v, ctx)``: ``F v``
    - ``factor_solve(v, ctx)``: ``F^{-1} v``
    - ``factor_solve_transpose(v, ctx)``: ``F^{-T} v``
    - ``norm(v, ctx)``: ``||v||_W``, defaulted to ``||factor_apply(v)||_2``.
      Provided for callers; the solvers measure in the whitened variable and
      never call it.

    Every op receives a :class:`~nlls_gram.SolverContext` carrying the
    solver's live state, so an exotic metric can key off the iterate through
    ``ctx.x`` and ``ctx.lm_state``.

    Metric instances are JAX PYTREES: array fields are traced leaves, the
    type plus its static fields are structure. Every concrete class must be
    registered with
    :func:`~nlls_gram.register_pytree_dataclass` -- the solvers reject
    unregistered instances. The instance rides inside the solver state
    (``lm_state.metric``), so a ``solve`` callback replaces the metric by
    constructing a new instance of the same type (same static fields, same
    leaf shapes and dtypes) -- pure traced ops, no recompilation. Equal-config
    instances with fresh arrays share one compiled solve loop.
    ``dataclasses.replace`` works too: metric constructors only validate and
    derive shapes, so re-running them under trace is cheap.

    ``free_scale`` weights the free block in the whitened variable: ``1.0``
    (the default) leaves it Euclidean. The ridge solver never penalizes the
    free block whatever the scale -- ``free_scale`` only changes its
    trust-region geometry -- while for the metric solver it IS that block's
    damping weight. It is a traced leaf, canonicalized by each constructor to
    the factor's float dtype, so changing it never recompiles.

    Contracts: the factor must be EXACT. The solver hardcodes the identity
    penalty block in the whitened variable, so an approximate factor silently
    changes the objective (unlike a CG preconditioner, which may be sloppy).
    The ridge weight never enters the factorization, so ridge continuation
    composes unchanged. How a subclass fulfills the ops -- prefactorized
    storage, factorize-in-``__init__``, fully matrix-free -- is its
    constructor's business, and the constructor must be traceable when the
    metric is rebuilt inside a jitted callback.
    """

    size: int
    free_scale: float = 1.0

    def factor_apply(self, v, ctx):
        """``F v`` for a metric-block vector or leading-axis-batched matrix."""
        raise NotImplementedError

    def factor_solve(self, v, ctx):
        """``F^{-1} v`` for a metric-block vector or batched matrix."""
        raise NotImplementedError

    def factor_solve_transpose(self, v, ctx):
        """``F^{-T} v`` for a metric-block vector or batched matrix."""
        raise NotImplementedError

    def norm(self, v, ctx):
        """``||v||_W = ||F v||_2`` for a metric-block vector."""
        return jnp.linalg.norm(self.factor_apply(v, ctx))

factor_apply(v, ctx)

F v for a metric-block vector or leading-axis-batched matrix.

Source code in src/nlls_gram/metrics.py

def factor_apply(self, v, ctx):
    """``F v`` for a metric-block vector or leading-axis-batched matrix."""
    raise NotImplementedError

factor_solve(v, ctx)

F^{-1} v for a metric-block vector or batched matrix.

Source code in src/nlls_gram/metrics.py

def factor_solve(self, v, ctx):
    """``F^{-1} v`` for a metric-block vector or batched matrix."""
    raise NotImplementedError

factor_solve_transpose(v, ctx)

F^{-T} v for a metric-block vector or batched matrix.

Source code in src/nlls_gram/metrics.py

def factor_solve_transpose(self, v, ctx):
    """``F^{-T} v`` for a metric-block vector or batched matrix."""
    raise NotImplementedError

norm(v, ctx)

||v||_W = ||F v||_2 for a metric-block vector.

Source code in src/nlls_gram/metrics.py

def norm(self, v, ctx):
    """``||v||_W = ||F v||_2`` for a metric-block vector."""
    return jnp.linalg.norm(self.factor_apply(v, ctx))

nlls_gram.IdentityMetric dataclass

Bases: Metric

The identity metric W = I on size coordinates -- plain ridge for the ridge solver, Euclidean damping for the metric solver.

F = I: every factor op is the identity and norm is the Euclidean norm, with no special-casing anywhere downstream.

Source code in src/nlls_gram/metrics.py

@dataclass(frozen=True, eq=False)
class IdentityMetric(Metric):
    """The identity metric ``W = I`` on ``size`` coordinates -- plain ridge for
    the ridge solver, Euclidean damping for the metric solver.

    ``F = I``: every factor op is the identity and ``norm`` is the Euclidean
    norm, with no special-casing anywhere downstream.
    """

    size: int
    free_scale: float = 1.0

    def __post_init__(self):
        if self.size < 0:
            raise ValueError("size must be nonnegative")
        object.__setattr__(
            self,
            "free_scale",
            _canonical_free_scale(self.free_scale, jnp.result_type(float)),
        )

    def factor_apply(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def factor_solve(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def factor_solve_transpose(self, v, ctx):
        _check_leading_size(v, self.size)
        return v

    def norm(self, v, ctx):
        _check_leading_size(v, self.size)
        return jnp.linalg.norm(v)

nlls_gram.CholeskyMetric dataclass

Bases: Metric

Dense metric W = L L' from its lower-triangular Cholesky factor.

L is the factor as jnp.linalg.cholesky returns it, so the upper factor this class works with is F = L'. Triangularity and positive definiteness are assumed, not validated -- the entries may be traced, and a singular factor propagates NaN loudly through the triangular solves.

Source code in src/nlls_gram/metrics.py

@dataclass(frozen=True, eq=False)
class CholeskyMetric(Metric):
    """Dense metric ``W = L L'`` from its lower-triangular Cholesky factor.

    ``L`` is the factor as ``jnp.linalg.cholesky`` returns it, so the upper
    factor this class works with is ``F = L'``. Triangularity and positive
    definiteness are assumed, not validated -- the entries may be traced, and
    a singular factor propagates NaN loudly through the triangular solves.
    """

    L: jax.Array
    free_scale: float = 1.0
    size: int = field(init=False)

    def __post_init__(self):
        L = jnp.asarray(self.L)
        if L.ndim != 2 or L.shape[0] != L.shape[1] or L.shape[0] == 0:
            raise ValueError("L must be a nonempty square matrix")
        object.__setattr__(self, "L", L)
        object.__setattr__(self, "size", L.shape[0])
        object.__setattr__(
            self,
            "free_scale",
            _canonical_free_scale(self.free_scale, jnp.result_type(L, 1.0)),
        )

    def factor_apply(self, v, ctx):
        _check_leading_size(v, self.size)
        return mm(self.L.T, v)

    def factor_solve(self, v, ctx):
        _check_leading_size(v, self.size)
        return jsp_linalg.solve_triangular(self.L.T, v, lower=False)

    def factor_solve_transpose(self, v, ctx):
        _check_leading_size(v, self.size)
        return jsp_linalg.solve_triangular(self.L, v, lower=True)

nlls_gram.DiagonalMetric dataclass

Bases: Metric

The diagonal metric W = diag(weights); F = diag(sqrt(weights)).

weights must be a positive 1-D array. Positivity is not validated because the values may be traced. Every op is elementwise.

Source code in src/nlls_gram/metrics.py

@dataclass(frozen=True, eq=False)
class DiagonalMetric(Metric):
    """The diagonal metric ``W = diag(weights)``; ``F = diag(sqrt(weights))``.

    ``weights`` must be a positive 1-D array. Positivity is not validated
    because the values may be traced. Every op is elementwise.
    """

    weights: jax.Array
    free_scale: float = 1.0
    size: int = field(init=False)

    def __post_init__(self):
        weights = jnp.asarray(self.weights)
        if weights.ndim != 1:
            raise ValueError("weights must be 1-D")
        object.__setattr__(self, "weights", weights)
        object.__setattr__(self, "size", weights.shape[0])
        object.__setattr__(
            self,
            "free_scale",
            _canonical_free_scale(self.free_scale, jnp.result_type(weights, 1.0)),
        )

    def _scaled(self, v, factor):
        _check_leading_size(v, self.size)
        return v * factor.reshape(factor.shape + (1,) * (v.ndim - 1))

    def factor_apply(self, v, ctx):
        return self._scaled(v, jnp.sqrt(self.weights))

    def factor_solve(self, v, ctx):
        return self._scaled(v, 1.0 / jnp.sqrt(self.weights))

    def factor_solve_transpose(self, v, ctx):
        return self.factor_solve(v, ctx)

nlls_gram.RepeatedFactorMetric dataclass

Bases: Metric

repeats copies of one block factor: W = blockdiag(F'F, ...).

F is an upper-triangular invertible block factor -- e.g. jnp.linalg.cholesky(K, upper=True) for a positive-definite Gram matrix K, giving the repeated kernel seminorm sum_j alpha_j' K alpha_j over repeats coefficient blocks. The constructor takes the FACTOR, not K (callers typically already hold it); a positive-SEMIdefinite K needs a shift first: jnp.linalg.cholesky(K + epsilon * I, upper=True). Triangularity and positive definiteness are assumed, not validated.

All repeated blocks (and all batched columns) share a single triangular product or solve: the ops reshape the metric block into the columns of one (block, repeats * cols) matrix, so no repeated factor or full block diagonal is ever formed. size = repeats * F.shape[0].

Source code in src/nlls_gram/metrics.py

@dataclass(frozen=True, eq=False)
class RepeatedFactorMetric(Metric):
    """``repeats`` copies of one block factor: ``W = blockdiag(F'F, ...)``.

    ``F`` is an upper-triangular invertible block factor -- e.g.
    ``jnp.linalg.cholesky(K, upper=True)`` for a positive-definite Gram matrix
    ``K``, giving the repeated kernel seminorm ``sum_j alpha_j' K alpha_j``
    over ``repeats`` coefficient blocks. The constructor takes the FACTOR, not
    ``K`` (callers typically already hold it); a positive-SEMIdefinite ``K``
    needs a shift first: ``jnp.linalg.cholesky(K + epsilon * I, upper=True)``.
    Triangularity and positive definiteness are assumed, not validated.

    All repeated blocks (and all batched columns) share a single triangular
    product or solve: the ops reshape the metric block into the columns of one
    ``(block, repeats * cols)`` matrix, so no repeated factor or full block
    diagonal is ever formed. ``size = repeats * F.shape[0]``.
    """

    F: jax.Array
    repeats: int = field(default=1, kw_only=True)
    free_scale: float = field(default=1.0, kw_only=True)
    size: int = field(init=False)

    def __post_init__(self):
        F = jnp.asarray(self.F)
        if F.ndim != 2 or F.shape[0] != F.shape[1] or F.shape[0] == 0:
            raise ValueError("F must be a nonempty square matrix")
        dtype = jnp.result_type(F, 1.0)
        if not jnp.issubdtype(dtype, jnp.floating):
            raise TypeError("F must have a real floating-point dtype")
        if self.repeats < 1:
            raise ValueError("repeats must be a positive integer")
        object.__setattr__(self, "F", F.astype(dtype))
        object.__setattr__(self, "size", self.repeats * F.shape[0])
        object.__setattr__(
            self, "free_scale", _canonical_free_scale(self.free_scale, dtype)
        )

    def _map_blocks(self, block_op, v):
        _check_leading_size(v, self.size)
        block_size = self.F.shape[0]
        trailing_shape = v.shape[1:]
        packed = jnp.moveaxis(
            v.reshape((self.repeats, block_size) + trailing_shape), 0, 1
        ).reshape(block_size, -1)
        return jnp.moveaxis(
            block_op(packed).reshape((block_size, self.repeats) + trailing_shape),
            0,
            1,
        ).reshape((self.size,) + trailing_shape)

    def factor_apply(self, v, ctx):
        return self._map_blocks(lambda m: mm(self.F, m), v)

    def factor_solve(self, v, ctx):
        return self._map_blocks(
            lambda m: jsp_linalg.solve_triangular(self.F, m, lower=False), v
        )

    def factor_solve_transpose(self, v, ctx):
        return self._map_blocks(
            lambda m: jsp_linalg.solve_triangular(self.F.T, m, lower=True), v
        )

Preconditioners

nlls_gram.Preconditioner

SPD preconditioner for the solvers' CG paths.

apply(v, damping, ctx) returns an SPD approximation of the damped operator's inverse -- (J~'J~ + ridge E + damping I)^{-1} in parameter space under :class:~nlls_gram.CG, (J~J~' + damping I)^{-1} in residual space under :class:~nlls_gram.GramCG. In the forward role it sits in CG's M slot with the live damping; in the ad_solver role the implicit system is undamped and damping is zero. ctx is the same :class:~nlls_gram.SolverContext the metric factor ops receive.

Preconditioner instances are JAX PYTREES: array fields are traced leaves, the type plus its static fields are structure. Every concrete class must be registered with :func:~nlls_gram.register_pytree_dataclass -- the solvers reject unregistered instances. The instance rides inside the solver state (lm_state.preconditioner), so a solve callback refreshes it by calling the CONSTRUCTOR again with fresh arrays -- same type, same leaf shapes and dtypes, pure traced ops, no recompilation. Never dataclasses.replace a preconditioner: replace re-runs __init__, re-paying any eigendecomposition or sketch, and construction-time-only inputs are not stored to re-supply. A stale instance only changes the CG iteration path, never the converged step, so refreshing rarely (or never) is always safe.

Implement a custom one as a small registered frozen dataclass::

@dataclass(frozen=True, eq=False)
class JacobiPreconditioner(Preconditioner):
    diagonal: jax.Array

    def apply(self, v, damping, ctx):
        return v / (self.diagonal + damping)

register_pytree_dataclass(JacobiPreconditioner, data_fields=("diagonal",))

Subclasses whose apply divides by the live damping must set requires_positive_damping = True; the constructor rejects them for the AD role, where damping is zero.

Source code in src/nlls_gram/preconditioners.py

class Preconditioner:
    """SPD preconditioner for the solvers' CG paths.

    ``apply(v, damping, ctx)`` returns an SPD approximation of the damped
    operator's inverse -- ``(J~'J~ + ridge E + damping I)^{-1}`` in parameter
    space under :class:`~nlls_gram.CG`, ``(J~J~' + damping I)^{-1}`` in
    residual space under :class:`~nlls_gram.GramCG`. In the forward role it
    sits in CG's ``M`` slot with the live damping; in the ``ad_solver`` role
    the implicit system is undamped and ``damping`` is zero. ``ctx`` is the
    same :class:`~nlls_gram.SolverContext` the metric factor ops receive.

    Preconditioner instances are JAX PYTREES: array fields are traced
    leaves, the type plus its static fields are structure. Every concrete
    class must be registered with
    :func:`~nlls_gram.register_pytree_dataclass` -- the solvers reject
    unregistered instances. The instance rides inside the solver state
    (``lm_state.preconditioner``), so a ``solve`` callback refreshes it by
    calling the CONSTRUCTOR again with fresh arrays -- same type, same leaf
    shapes and dtypes, pure traced ops, no recompilation. Never
    ``dataclasses.replace`` a preconditioner: ``replace`` re-runs
    ``__init__``, re-paying any eigendecomposition or sketch, and
    construction-time-only inputs are not stored to re-supply. A stale
    instance only changes the CG iteration path, never the converged step,
    so refreshing rarely (or never) is always safe.

    Implement a custom one as a small registered frozen dataclass::

        @dataclass(frozen=True, eq=False)
        class JacobiPreconditioner(Preconditioner):
            diagonal: jax.Array

            def apply(self, v, damping, ctx):
                return v / (self.diagonal + damping)

        register_pytree_dataclass(JacobiPreconditioner, data_fields=("diagonal",))

    Subclasses whose ``apply`` divides by the live damping must set
    ``requires_positive_damping = True``; the constructor rejects them for
    the AD role, where damping is zero.
    """

    requires_positive_damping = False

    def apply(self, v, damping, ctx):
        raise NotImplementedError

nlls_gram.IdentityPreconditioner dataclass

Bases: Preconditioner

The identity map as an explicit "no preconditioner" choice for CG.

Nobody should run Krylov methods without thinking about preconditioning, so opting out is an explicit, greppable decision rather than a silent default. Stateless: every instance compares equal, so equal CG configs share one compiled solve loop.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True)
class IdentityPreconditioner(Preconditioner):
    """The identity map as an explicit "no preconditioner" choice for ``CG``.

    Nobody should run Krylov methods without thinking about preconditioning,
    so opting out is an explicit, greppable decision rather than a silent
    default. Stateless: every instance compares equal, so equal ``CG``
    configs share one compiled solve loop.
    """

    def apply(self, v, damping, ctx):
        return v

nlls_gram.BlockEigenPreconditioner dataclass

Bases: Preconditioner

Block-diagonal eigenbasis preconditioner over grouped whitened coordinates.

The workhorse for structured whitened normal operators J~'J~ + ridge E + damping I built from repeated interacting blocks (multiple "agents" coupled through shared equations): approximate the operator by a block-diagonal matrix over a chosen grouping of the whitened coordinates, eigendecompose each block, and apply the exact inverse of the shifted approximation

v  ->  V ((V' v) / (Lambda + ridge_weight * ridge + damping)) V'

per block -- analytic in both the live damping (traced; it changes per LM step) and the live ridge (read from ctx.lm_state.ridge, so ridge continuation composes with no refresh).

families is a sequence of (blocks, ridge_weight) pairs: blocks has shape (groups, size, size) -- the stacked diagonal blocks of J~'J~ restricted to that family's coordinate groups, in permuted order. Families in the metric block set ridge_weight = 1 (their diagonal carries the ridge spectral floor); free-block families set 0 (damping-only, so the zero-damping AD role applies their plain inverse -- positive definite whenever the free block is identified). Blocks are symmetrized and eigendecomposed HERE, once; positive semidefiniteness is assumed, not validated (entries may be traced). permutation is the 1-D integer array reordering the flattened whitened parameter vector into family-major order (v_permuted = v[permutation]); families are consumed in sequence and must cover it exactly. jnp.arange(n) serves when the natural layout already is family-major.

The constructor is fully traceable, so a solve callback refreshes the preconditioner from the live iterate by constructing a new instance -- gate the rebuild with jax.lax.cond (e.g. on a ridge-continuation level advance) so the eigendecomposition is paid only when it fires; a jnp.where merge would pay it every step.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True, eq=False)
class BlockEigenPreconditioner(Preconditioner):
    """Block-diagonal eigenbasis preconditioner over grouped whitened
    coordinates.

    The workhorse for structured whitened normal operators
    ``J~'J~ + ridge E + damping I`` built from repeated interacting blocks
    (multiple "agents" coupled through shared equations): approximate the
    operator by a block-diagonal matrix over a chosen grouping of the
    whitened coordinates, eigendecompose each block, and apply the exact
    inverse of the shifted approximation

        v  ->  V ((V' v) / (Lambda + ridge_weight * ridge + damping)) V'

    per block -- analytic in both the live ``damping`` (traced; it changes per
    LM step) and the live ``ridge`` (read from ``ctx.lm_state.ridge``, so
    ridge continuation composes with no refresh).

    ``families`` is a sequence of ``(blocks, ridge_weight)`` pairs:
    ``blocks`` has shape ``(groups, size, size)`` -- the stacked diagonal
    blocks of ``J~'J~`` restricted to that family's coordinate groups, in
    permuted order. Families in the metric block set ``ridge_weight = 1``
    (their diagonal carries the ``ridge`` spectral floor); free-block
    families set ``0`` (damping-only, so the zero-damping AD role applies
    their plain inverse -- positive definite whenever the free block is
    identified). Blocks are symmetrized and eigendecomposed HERE, once;
    positive semidefiniteness is assumed, not validated (entries may be
    traced). ``permutation`` is the 1-D integer array reordering the
    flattened whitened parameter vector into family-major order
    (``v_permuted = v[permutation]``); families are consumed in sequence and
    must cover it exactly. ``jnp.arange(n)`` serves when the natural layout
    already is family-major.

    The constructor is fully traceable, so a ``solve`` callback refreshes
    the preconditioner from the live iterate by constructing a new instance
    -- gate the rebuild with ``jax.lax.cond`` (e.g. on a ridge-continuation
    level advance) so the eigendecomposition is paid only when it fires; a
    ``jnp.where`` merge would pay it every step.
    """

    families: InitVar[Any]
    permutation: jax.Array
    eigenvectors: tuple = field(init=False)
    eigenvalues: tuple = field(init=False)
    ridge_weights: tuple = field(init=False)
    inverse_permutation: jax.Array = field(init=False)

    def __post_init__(self, families):
        permutation = jnp.asarray(self.permutation)
        if permutation.ndim != 1 or not jnp.issubdtype(permutation.dtype, jnp.integer):
            raise ValueError("permutation must be a 1-D integer array")
        eigenvectors, eigenvalues, ridge_weights = [], [], []
        covered = 0
        for blocks, ridge_weight in families:
            blocks = jnp.asarray(blocks)
            if blocks.ndim != 3 or blocks.shape[1] != blocks.shape[2]:
                raise ValueError(
                    "each family's blocks must have shape (groups, size, size); "
                    f"got {blocks.shape}"
                )
            symmetrized = 0.5 * (blocks + jnp.swapaxes(blocks, 1, 2))
            values, vectors = jnp.linalg.eigh(symmetrized)
            # eigh of a numerically PSD block can return tiny negative
            # eigenvalues; clamped at zero the apply shift stays positive for
            # any positive ridge/damping (and the zero-damping AD role stays
            # SPD whenever the family itself is).
            eigenvalues.append(jnp.maximum(values, 0.0))
            eigenvectors.append(vectors)
            ridge_weights.append(jnp.asarray(ridge_weight, dtype=blocks.dtype))
            covered += blocks.shape[0] * blocks.shape[1]
        if covered != permutation.shape[0]:
            raise ValueError(
                f"families cover {covered} coordinates but the permutation "
                f"has {permutation.shape[0]}"
            )
        object.__setattr__(self, "permutation", permutation)
        object.__setattr__(self, "eigenvectors", tuple(eigenvectors))
        object.__setattr__(self, "eigenvalues", tuple(eigenvalues))
        object.__setattr__(self, "ridge_weights", tuple(ridge_weights))
        object.__setattr__(self, "inverse_permutation", jnp.argsort(permutation))

    def apply(self, v, damping, ctx):
        # LevenbergMarquardt carries no ridge, so its metric-block families
        # shift by the damping alone.
        carried = ctx.lm_state.ridge
        ridge = (
            jnp.zeros((), v.dtype)
            if carried is None
            else jnp.asarray(carried, dtype=v.dtype)
        )
        permuted = v[self.permutation]
        pieces = []
        offset = 0
        for V, values, ridge_weight in zip(
            self.eigenvectors, self.eigenvalues, self.ridge_weights, strict=True
        ):
            groups, size = V.shape[0], V.shape[1]
            segment = permuted[offset : offset + groups * size]
            offset += groups * size
            shift = ridge_weight * ridge + damping
            coefficients = jnp.einsum(
                "gab,ga->gb", V, segment.reshape(groups, size), precision=HIGHEST
            )
            pieces.append(
                jnp.einsum(
                    "gab,gb->ga", V, coefficients / (values + shift), precision=HIGHEST
                ).reshape(-1)
            )
        return jnp.concatenate(pieces)[self.inverse_permutation].astype(v.dtype)

nlls_gram.NystromPreconditioner dataclass

Bases: Preconditioner

Randomized Nystrom preconditioner (Frangella-Tropp-Udell) for a PSD operator given only through matvec.

Sketches A with a rank-rank Nystrom approximation A_hat = U diag(lam) U' -- a thin-QR'd Gaussian test matrix, one block application Y = A Omega, and the shifted Cholesky/SVD recovery of arXiv:2110.02820 Algorithm 2.1 -- then applies their eq. 5.3::

v  ->  U ((U'v) / (lam + damping)) + (v - U U'v) / (rho + damping)

where rho is the smallest retained eigenvalue: directions the sketch resolved are inverted against the live shift, and the unresolved complement is treated as sitting at rho rather than at zero. That balance is what carries the FTU condition-number guarantee for fast-decaying spectra.

The target use is neural-network least squares under the identity metric, where the dual operator is the m x m empirical NTK Gram J J' -- fast spectral decay plus the LM damping shift is exactly the FTU regime. matvec must apply a symmetric PSD operator to (n, k) matrices; an indefinite one silently produces NaN through the Cholesky square root. It is consumed at construction -- the build costs rank operator applications plus an O(n rank^2) QR/SVD, and only the sketch is stored, so for a nonlinear problem the instance approximates the dual at the linearization point it was built from (staleness is safe; a callback refreshes by constructing a new instance from a fresh matvec). Each apply is two (n, rank) matvecs.

key is an explicit PRNG key; the same key reproduces the same preconditioner. dtype=None uses the JAX default float -- pass the operator dtype explicitly for a float32 problem under enabled x64.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True, eq=False)
class NystromPreconditioner(Preconditioner):
    """Randomized Nystrom preconditioner (Frangella-Tropp-Udell) for a PSD
    operator given only through ``matvec``.

    Sketches ``A`` with a rank-``rank`` Nystrom approximation
    ``A_hat = U diag(lam) U'`` -- a thin-QR'd Gaussian test matrix, one block
    application ``Y = A Omega``, and the shifted Cholesky/SVD recovery of
    arXiv:2110.02820 Algorithm 2.1 -- then applies their eq. 5.3::

        v  ->  U ((U'v) / (lam + damping)) + (v - U U'v) / (rho + damping)

    where ``rho`` is the smallest retained eigenvalue: directions the sketch
    resolved are inverted against the live shift, and the unresolved
    complement is treated as sitting at ``rho`` rather than at zero. That
    balance is what carries the FTU condition-number guarantee for
    fast-decaying spectra.

    The target use is neural-network least squares under the identity metric,
    where the dual operator is the ``m x m`` empirical NTK Gram ``J J'`` --
    fast spectral decay plus the LM damping shift is exactly the FTU regime.
    ``matvec`` must apply a symmetric PSD operator to ``(n, k)`` matrices; an
    indefinite one silently produces NaN through the Cholesky square root. It
    is consumed at construction -- the build costs ``rank`` operator
    applications plus an ``O(n rank^2)`` QR/SVD, and only the sketch is
    stored, so for a nonlinear problem the instance approximates the dual at
    the linearization point it was built from (staleness is safe; a callback
    refreshes by constructing a new instance from a fresh ``matvec``). Each
    apply is two ``(n, rank)`` matvecs.

    ``key`` is an explicit PRNG key; the same key reproduces the same
    preconditioner. ``dtype=None`` uses the JAX default float -- pass the
    operator dtype explicitly for a float32 problem under enabled x64.
    """

    matvec: InitVar[Any]
    n: int
    rank: int
    key: InitVar[Any]
    dtype: InitVar[Any] = None
    basis: jax.Array = field(init=False)
    eigenvalues: jax.Array = field(init=False)

    def __post_init__(self, matvec, key, dtype):
        if not 0 < self.rank <= self.n:
            raise ValueError("rank must be a positive int <= n")
        dtype = jnp.result_type(float) if dtype is None else dtype
        shape = (self.n, self.rank)
        Omega = jnp.linalg.qr(jax.random.normal(key, shape, dtype))[0]
        Y = matvec(Omega)
        # The floor keeps the shift usable for a (near-)zero operator, where
        # eps * ||Y||_F alone would leave the core singular; tiny/eps stays
        # clear of the subnormal range through the downstream products.
        finfo = jnp.finfo(dtype)
        nu = jnp.maximum(finfo.eps * jnp.linalg.norm(Y), finfo.tiny / finfo.eps)
        Y_nu = Y + nu * Omega
        core = mm(Omega.T, Y_nu)
        L = jnp.linalg.cholesky(0.5 * (core + core.T))
        B = jsp_linalg.solve_triangular(L, Y_nu.T, lower=True).T
        U, sigma, _ = jnp.linalg.svd(B, full_matrices=False)
        object.__setattr__(self, "basis", U)
        object.__setattr__(self, "eigenvalues", jnp.maximum(sigma**2 - nu, 0.0))

    def apply(self, v, damping, ctx):
        # Regrouped so the apply is two (n, rank) matvecs instead of three.
        U, lam = self.basis, self.eigenvalues
        rho = lam[-1]
        Utv = mm(U.T, v)
        return mm(U, Utv / (lam + damping) - Utv / (rho + damping)) + v / (
            rho + damping
        )

nlls_gram.ShermanMorrisonPreconditioner dataclass

Bases: Preconditioner

Dual preconditioner for B = A + weight * u u' from a solve with A.

Applies B^{-1} v = y - A^{-1}u (u' y) / (1/weight + u' A^{-1} u) with y = A^{-1} v by the Sherman-Morrison identity; A^{-1}u and the scalar denominator are precomputed at construction. This is the natural shape for kernel-collocation dual operators, where a metric weight m on a scalar parameter injects an exactly known rank-1 spike (c^2/m) u u' into J M^{-1} J'. The live damping is ignored -- spectral closeness to the damped operator is all a preconditioner needs -- which also makes it valid in the zero-damping ad_solver role.

solve applies A^{-1} and is called on every apply, so it is a STATIC field: a fixed hashable callable whose identity enters the instance's pytree structure, with anything it closes over entering the compiled program as constants. This class is a setup-scope object for a fixed dual operator, not a callback-refresh target.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True, eq=False)
class ShermanMorrisonPreconditioner(Preconditioner):
    """Dual preconditioner for ``B = A + weight * u u'`` from a solve with ``A``.

    Applies ``B^{-1} v = y - A^{-1}u (u' y) / (1/weight + u' A^{-1} u)`` with
    ``y = A^{-1} v`` by the Sherman-Morrison identity; ``A^{-1}u`` and the
    scalar denominator are precomputed at construction. This is the natural
    shape for kernel-collocation dual operators, where a metric weight ``m``
    on a scalar parameter injects an exactly known rank-1 spike
    ``(c^2/m) u u'`` into ``J M^{-1} J'``. The live ``damping`` is ignored --
    spectral closeness to the damped operator is all a preconditioner needs --
    which also makes it valid in the zero-damping ``ad_solver`` role.

    ``solve`` applies ``A^{-1}`` and is called on every ``apply``, so it is a
    STATIC field: a fixed hashable callable whose identity enters the
    instance's pytree structure, with anything it closes over entering the
    compiled program as constants. This class is a setup-scope object for a
    fixed dual operator, not a callback-refresh target.
    """

    solve: Any
    u: jax.Array
    weight: Any
    solve_u: jax.Array = field(init=False)
    denominator: jax.Array = field(init=False)

    def __post_init__(self):
        u = jnp.asarray(self.u)
        weight = jnp.asarray(self.weight, dtype=jnp.result_type(u, 1.0))
        solve_u = self.solve(u)
        object.__setattr__(self, "u", u)
        object.__setattr__(self, "weight", weight)
        object.__setattr__(self, "solve_u", solve_u)
        object.__setattr__(self, "denominator", 1.0 / weight + mm(u, solve_u))

    def apply(self, v, damping, ctx):
        y = self.solve(v)
        return y - self.solve_u * (mm(self.u, y) / self.denominator)

nlls_gram.WoodburyPreconditioner dataclass

Bases: Preconditioner

Dual preconditioner for B = A + U diag(weights) U'.

The rank-k generalization of :class:ShermanMorrisonPreconditioner: applies B^{-1} v = y - A^{-1}U C^{-1}(U' y) with y = A^{-1} v and capacitance C = diag(1/weights) + U' A^{-1} U; A^{-1}U (one matrix solve) and the Cholesky factor of the k x k capacitance are precomputed. weights must be positive -- not validated, since inputs may be traced. Like Sherman-Morrison it ignores damping and so serves the AD role too, and its solve is the same STATIC always-called field.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True, eq=False)
class WoodburyPreconditioner(Preconditioner):
    """Dual preconditioner for ``B = A + U diag(weights) U'``.

    The rank-k generalization of :class:`ShermanMorrisonPreconditioner`:
    applies ``B^{-1} v = y - A^{-1}U C^{-1}(U' y)`` with ``y = A^{-1} v`` and
    capacitance ``C = diag(1/weights) + U' A^{-1} U``; ``A^{-1}U`` (one matrix
    solve) and the Cholesky factor of the k x k capacitance are precomputed.
    ``weights`` must be positive -- not validated, since inputs may be traced.
    Like Sherman-Morrison it ignores ``damping`` and so serves the AD role
    too, and its ``solve`` is the same STATIC always-called field.
    """

    solve: Any
    U: jax.Array
    weights: jax.Array
    solve_U: jax.Array = field(init=False)
    capacitance_factor: jax.Array = field(init=False)

    def __post_init__(self):
        U, weights = jnp.asarray(self.U), jnp.asarray(self.weights)
        if U.ndim != 2 or weights.shape != (U.shape[1],):
            raise ValueError("U must have shape (n, k) and weights shape (k,)")
        object.__setattr__(self, "U", U)
        object.__setattr__(self, "weights", weights)
        solve_U = self.solve(U)
        object.__setattr__(self, "solve_U", solve_U)
        capacitance = jnp.diag(1.0 / weights) + mm(U.T, solve_U)
        object.__setattr__(
            self, "capacitance_factor", jsp_linalg.cho_factor(capacitance)[0]
        )

    def apply(self, v, damping, ctx):
        y = self.solve(v)
        correction = jsp_linalg.cho_solve(
            (self.capacitance_factor, False), mm(self.U.T, y)
        )
        return y - mm(self.solve_U, correction)

nlls_gram.PaddedPreconditioner dataclass

Bases: Preconditioner

Extend a dual preconditioner to a residual padded with exact zeros.

The fixed-residual-shape pattern appends identically-zero entries to an n_real-entry residual so compiled shapes stay stable across problem instances. Padded rows have zero Jacobian rows, so the dual operator is exactly block diagonal::

[ J P J' + damping I      0          ]
[ 0                       damping I  ]

and this applies base on the first n_real coordinates and the exact 1/damping inverse on the padded block. That second block must NOT be zeroed -- that would make the preconditioner singular rather than SPD, even though zeros can appear to work when the padded coordinates are never excited. Because the padded block divides by the live damping, this serves only the damped forward solve; the undamped dual is singular there, which ad_solver=SVD() handles exactly.

Source code in src/nlls_gram/preconditioners.py

@dataclass(frozen=True, eq=False)
class PaddedPreconditioner(Preconditioner):
    """Extend a dual preconditioner to a residual padded with exact zeros.

    The fixed-residual-shape pattern appends identically-zero entries to an
    ``n_real``-entry residual so compiled shapes stay stable across problem
    instances. Padded rows have zero Jacobian rows, so the dual operator is
    exactly block diagonal::

        [ J P J' + damping I      0          ]
        [ 0                       damping I  ]

    and this applies ``base`` on the first ``n_real`` coordinates and the exact
    ``1/damping`` inverse on the padded block. That second block must NOT be
    zeroed -- that would make the preconditioner singular rather than SPD, even
    though zeros can appear to work when the padded coordinates are never
    excited. Because the padded block divides by the live damping, this serves
    only the damped forward solve; the undamped dual is singular there, which
    ``ad_solver=SVD()`` handles exactly.
    """

    base: Preconditioner
    n_real: int

    requires_positive_damping = True

    def apply(self, v, damping, ctx):
        # Static shapes, so this raises at trace time; without it a
        # shape-generic base would silently accept a too-short vector.
        if v.ndim != 1 or v.shape[0] < self.n_real:
            raise ValueError(
                f"padded residual vector must be 1-D with at least "
                f"n_real={self.n_real} entries; got shape {v.shape}"
            )
        return jnp.concatenate(
            (
                self.base.apply(v[: self.n_real], damping, ctx),
                v[self.n_real :] / damping,
            )
        )

State and results

nlls_gram.LMState dataclass

Carried solver state threaded through init/update/solve.

Only damping is always live; every other field is populated by the configuration that needs it and stays None otherwise. A solve callback that rebuilds the state must PRESERVE the fields it does not mean to change -- use dataclasses.replace(ctx.lm_state, ...).

The callback-owned fields are damping, ridge, metric, preconditioner, and the hyper group (replaceable as a unit); everything else is solver-owned bookkeeping.

Attributes:

Name	Type	Description
damping	Array	() current LM damping.
ridge	Array | None	() ridge weight, strictly positive, for RidgeLevenbergMarquardt; None for the metric solver. Replacing it is the supported way to anneal mid-solve (see AnnealRidge); the solver treats a ridge change as a problem change, suppressing that step's convergence test and invalidating the ridge-keyed caches.
resid	Array | None	cached residual at the current x (cache_jacobian dense paths only).
Jt	Array | None	cached transpose-Jacobian J' at the current x.
jacobian_valid	Array | None	() bool -- the cached resid/Jt are still current because the last step was rejected, so x did not move.
aux	Any	residual aux pytree at the current x (has_aux=True).
hyper	LMHyperparams | None	per-step :class:LMHyperparams, populated by solve; None (init's default) falls back to the constructor values.
solver_cache	Any	the linear solver's own reject-step cache, whose pytree structure is fixed by the static linear_solver config.
metric	Any	the carried :class:~nlls_gram.Metric instance every factor op reads. A callback replaces it by constructing a new instance of the same type with matching leaf shapes and dtypes; the ridge solver treats a changed metric as a problem change (the metric defines its objective), the metric solver only invalidates the whitening-dependent solver caches.
preconditioner	Any	the carried :class:~nlls_gram.Preconditioner instance the Krylov configs read; None outside them. A callback refresh is never treated as a problem change -- staleness only moves the CG iteration path.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMState:
    """Carried solver state threaded through ``init``/``update``/``solve``.

    Only ``damping`` is always live; every other field is populated by the
    configuration that needs it and stays ``None`` otherwise. A ``solve``
    callback that rebuilds the state must PRESERVE the fields it does not mean
    to change -- use ``dataclasses.replace(ctx.lm_state, ...)``.

    The callback-owned fields are ``damping``, ``ridge``, ``metric``,
    ``preconditioner``, and the ``hyper`` group (replaceable as a unit);
    everything else is solver-owned bookkeeping.

    Attributes:
        damping: ``()`` current LM damping.
        ridge: ``()`` ridge weight, strictly positive, for
            ``RidgeLevenbergMarquardt``; ``None`` for the metric solver.
            Replacing it is the supported way to anneal mid-solve (see
            ``AnnealRidge``); the solver treats a ridge change as a problem
            change, suppressing that step's convergence test and
            invalidating the ridge-keyed caches.
        resid: cached residual at the current ``x`` (``cache_jacobian`` dense
            paths only).
        Jt: cached transpose-Jacobian ``J'`` at the current ``x``.
        jacobian_valid: ``()`` bool -- the cached ``resid``/``Jt`` are still
            current because the last step was rejected, so ``x`` did not move.
        aux: residual aux pytree at the current ``x`` (``has_aux=True``).
        hyper: per-step :class:`LMHyperparams`, populated by ``solve``;
            ``None`` (``init``'s default) falls back to the constructor values.
        solver_cache: the linear solver's own reject-step cache, whose pytree
            structure is fixed by the static ``linear_solver`` config.
        metric: the carried :class:`~nlls_gram.Metric` instance every factor
            op reads. A callback replaces it by constructing a new instance
            of the same type with matching leaf shapes and dtypes; the ridge
            solver treats a changed metric as a problem change (the metric
            defines its objective), the metric solver only invalidates the
            whitening-dependent solver caches.
        preconditioner: the carried :class:`~nlls_gram.Preconditioner`
            instance the Krylov configs read; ``None`` outside them. A
            callback refresh is never treated as a problem change --
            staleness only moves the CG iteration path.
    """

    damping: jax.Array
    ridge: jax.Array | None = None
    resid: jax.Array | None = None
    Jt: jax.Array | None = None
    jacobian_valid: jax.Array | None = None
    aux: Any = None
    hyper: LMHyperparams | None = None
    solver_cache: Any = None
    metric: Any = None
    preconditioner: Any = None

nlls_gram.LMInfo dataclass

Per-step diagnostics returned by update and by each solve step.

The loss/damping fields report the accept/reject outcome of the step, while grad_norm/step_norm/aux are evaluated at the PRE-step x -- the iterate the step was computed from.

loss is always the objective the running solver minimizes: the sum of squared residuals for LevenbergMarquardt, and the RIDGE OBJECTIVE ||r||^2 + ridge * ||x_m||_W^2 (penalty included) for RidgeLevenbergMarquardt -- ridge code that means equation error must read resid_loss.

BOTH solvers run in the whitened variable y = F_bar x, so grad_norm, step_norm, and penalty_grad_norm are Euclidean in y: steps measured in the W-norm, gradients in the dual W^{-1}-norm. With the default Euclidean metric they are the plain quantities. Objective values are unaffected -- whitening is a linear bijection of the same objective.

Attributes:

Name	Type	Description
loss	Array	objective at the retained iterate, min(loss_old, loss_candidate).
loss_old	Array	objective at the pre-step x.
loss_candidate	Array	objective at the trial point.
accepted	Array	() bool, whether the trial step was accepted.
damping	Array	() post-update damping.
damping_factor	Array	() multiplicative damping update applied this step.
used_geodesic	Array	() bool, whether the geodesic correction entered the accepted step.
acceleration_ratio	Array	() acceleration-to-velocity norm ratio.
grad_norm	Array	() whitened stationarity residual at the pre-step x: ||F_bar^{-T} J'r||, plus ridge [y_m; 0] for the ridge solver.
step_norm	Array	() norm of the candidate step, reported even when the step is rejected.
ridge	Array | None	() the ridge weight used this step (ridge solver only).
resid_loss	Array | None	||r||^2 at the retained iterate (ridge solver only).
penalty_value	Array | None	||x_m||_W^2 = ||y_m||^2 at the retained iterate.
penalty_grad_norm	Array | None	() ||[y_m; 0]|| = sqrt(penalty_value) at the pre-step x, reported so gtol can be CALIBRATED rather than guessed: at a ridge minimizer the gradient is the cancellation of the residual pullback against ridge * [y_m; 0], so demanding grad_norm < c * ridge * penalty_grad_norm resolves the selection coordinates to ~c relative accuracy. The recipe is gtol ~ 1e-3 * ridge * sqrt(q(x*)) with q the solution's squared seminorm.
aux	Any	residual aux output at the pre-step x (has_aux=True).

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMInfo:
    """Per-step diagnostics returned by ``update`` and by each ``solve`` step.

    The loss/damping fields report the accept/reject outcome of the step, while
    ``grad_norm``/``step_norm``/``aux`` are evaluated at the PRE-step ``x`` --
    the iterate the step was computed from.

    ``loss`` is always the objective the running solver minimizes: the sum of
    squared residuals for ``LevenbergMarquardt``, and the RIDGE OBJECTIVE
    ``||r||^2 + ridge * ||x_m||_W^2`` (penalty included) for
    ``RidgeLevenbergMarquardt`` -- ridge code that means equation error must
    read ``resid_loss``.

    BOTH solvers run in the whitened variable ``y = F_bar x``, so
    ``grad_norm``, ``step_norm``, and ``penalty_grad_norm`` are Euclidean in
    ``y``: steps measured in the W-norm, gradients in the dual W^{-1}-norm.
    With the default Euclidean metric they are the plain quantities. Objective
    values are unaffected -- whitening is a linear bijection of the same
    objective.

    Attributes:
        loss: objective at the retained iterate, ``min(loss_old,
            loss_candidate)``.
        loss_old: objective at the pre-step ``x``.
        loss_candidate: objective at the trial point.
        accepted: ``()`` bool, whether the trial step was accepted.
        damping: ``()`` post-update damping.
        damping_factor: ``()`` multiplicative damping update applied this step.
        used_geodesic: ``()`` bool, whether the geodesic correction entered the
            accepted step.
        acceleration_ratio: ``()`` acceleration-to-velocity norm ratio.
        grad_norm: ``()`` whitened stationarity residual at the pre-step
            ``x``: ``||F_bar^{-T} J'r||``, plus ``ridge [y_m; 0]`` for the
            ridge solver.
        step_norm: ``()`` norm of the candidate step, reported even when the
            step is rejected.
        ridge: ``()`` the ridge weight used this step (ridge solver only).
        resid_loss: ``||r||^2`` at the retained iterate (ridge solver only).
        penalty_value: ``||x_m||_W^2 = ||y_m||^2`` at the retained iterate.
        penalty_grad_norm: ``()`` ``||[y_m; 0]|| = sqrt(penalty_value)`` at the
            pre-step ``x``, reported so ``gtol`` can be CALIBRATED rather than
            guessed: at a ridge minimizer the gradient is the cancellation of
            the residual pullback against ``ridge * [y_m; 0]``, so demanding
            ``grad_norm < c * ridge * penalty_grad_norm`` resolves the
            selection coordinates to ~``c`` relative accuracy. The recipe is
            ``gtol ~ 1e-3 * ridge * sqrt(q(x*))`` with ``q`` the solution's
            squared seminorm.
        aux: residual aux output at the pre-step ``x`` (``has_aux=True``).
    """

    loss: jax.Array
    loss_old: jax.Array
    loss_candidate: jax.Array
    accepted: jax.Array
    damping: jax.Array
    damping_factor: jax.Array
    used_geodesic: jax.Array
    acceleration_ratio: jax.Array
    grad_norm: jax.Array
    step_norm: jax.Array
    ridge: jax.Array | None = None
    resid_loss: jax.Array | None = None
    penalty_value: jax.Array | None = None
    penalty_grad_norm: jax.Array | None = None
    aux: Any = None

nlls_gram.LMStatus

Bases: IntEnum

Integer status codes returned by solve.

Members are real ints (IntEnum): they work as dict keys, compare against status arrays, and LMStatus(int(result.status)).name recovers the label for logging. Callbacks may return bare members (or any weak integer value) as LMAction.status -- the solver canonicalizes to int32 at the boundary, so no explicit dtype casts are needed.

Source code in src/nlls_gram/lm_types.py

class LMStatus(enum.IntEnum):
    """Integer status codes returned by ``solve``.

    Members are real ints (``IntEnum``): they work as dict keys, compare
    against status arrays, and ``LMStatus(int(result.status)).name`` recovers
    the label for logging. Callbacks may return bare members (or any weak
    integer value) as ``LMAction.status`` -- the solver canonicalizes to
    int32 at the boundary, so no explicit dtype casts are needed.
    """

    RUNNING = 0
    CONVERGED = 1
    MAX_STEPS = 2
    NONFINITE = 3
    CALLBACK_STOP = 4

nlls_gram.LMHyperparams dataclass

Per-step LM hyperparameters, carried in LMState.hyper.

All fields are traced, so a solve callback can reset them -- e.g. grow the inner CG budget as the loss falls -- via dataclasses.replace(ctx.lm_state, hyper=dataclasses.replace( ctx.lm_state.hyper, iterative_maxiter=...)). A field constructed as None (backend-default iterative_maxiter) is compiled out and stays None. Static configuration -- the linear solver, the metric, geodesic_acceleration, cache_jacobian, has_aux -- shapes the compiled program and lives on the solver.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMHyperparams:
    """Per-step LM hyperparameters, carried in ``LMState.hyper``.

    All fields are traced, so a ``solve`` callback can reset them -- e.g. grow
    the inner CG budget as the loss falls -- via
    ``dataclasses.replace(ctx.lm_state, hyper=dataclasses.replace(
    ctx.lm_state.hyper, iterative_maxiter=...))``. A field constructed as
    ``None`` (backend-default ``iterative_maxiter``) is compiled out and stays
    ``None``. Static configuration -- the linear solver, the metric,
    ``geodesic_acceleration``, ``cache_jacobian``, ``has_aux`` -- shapes the
    compiled program and lives on the solver.
    """

    damping_decrease: jax.Array
    damping_increase: jax.Array
    geodesic_acceptance_ratio: jax.Array
    iterative_tol: jax.Array
    iterative_atol: jax.Array
    iterative_maxiter: jax.Array | None

nlls_gram.LMContext dataclass

Information passed to a solve callback after each LM update.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMContext:
    """Information passed to a ``solve`` callback after each LM update."""

    step: jax.Array
    x: Any
    x_old: Any
    lm_state: Any
    lm_state_old: Any
    initial_lm_state: Any
    args: Any
    p: Any
    user_state: Any
    info: Any

nlls_gram.LMAction dataclass

Optional callback action for solve.

A field left as None is unchanged. status is used only when stop is true. stop and status are canonicalized by the solver (to bool and int32), so callbacks may return Python bools, bare LMStatus members, or weak-typed arrays without explicit casts.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMAction:
    """Optional callback action for ``solve``.

    A field left as ``None`` is unchanged. ``status`` is used only when
    ``stop`` is true. ``stop`` and ``status`` are canonicalized by the solver
    (to bool and int32), so callbacks may return Python bools, bare
    ``LMStatus`` members, or weak-typed arrays without explicit casts.
    """

    stop: Any = None
    status: Any = None
    x: Any = None
    lm_state: Any = None
    args: Any = None
    user_state: Any = None

nlls_gram.LMSolveResult dataclass

Final result returned by solve.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class LMSolveResult:
    """Final result returned by ``solve``."""

    x: Any
    lm_state: Any
    info: Any
    steps: jax.Array
    status: jax.Array
    args: Any
    p: Any
    user_state: Any
    # With has_aux=True: aux evaluated at the returned (x, args, p) -- one extra
    # residual evaluation, well-defined for every status. Differentiable with
    # respect to p through the implicit rule (directly and through x*(p)).
    aux: Any = None
    # With save_steps=True: the iterate history as a pytree shaped like x with a
    # (max_steps + 1) leading axis -- row 0 is x0, row s the kept iterate after
    # step s (post-callback-action), rows beyond ``steps`` are zero padding.
    # aux_history (has_aux only) and args_history (None when args is None) align
    # row-for-row. Differentiation-inert (zero tangents through the implicit rule).
    x_history: Any = None
    aux_history: Any = None
    args_history: Any = None
    # MultiStartInfo when solve ran with multi_start=...; None otherwise (an
    # empty pytree node, so the leaf count is unchanged when the feature is off).
    multi_start: Any = None

nlls_gram.SolverContext dataclass

What the solver knows at a metric, preconditioner, or linear-solver call site -- the inner algebra's context, as opposed to :class:LMContext, which a per-step user callback receives.

Fields are None where the call site has nothing to offer:

• x: the current FLATTENED iterate (the whole parameter vector, not just the metric block).
• lm_state: the live :class:LMState (damping, ridge, caches, and the carried metric/preconditioner instances). In the implicit-AD rule this is the returned state under stop_gradient -- inert conditioning data, like the ridge.
• args / p: the residual's auxiliary data and differentiation parameters as passed to solve/update.

Source code in src/nlls_gram/lm_types.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SolverContext:
    """What the solver knows at a metric, preconditioner, or linear-solver
    call site -- the inner algebra's context, as opposed to
    :class:`LMContext`, which a per-step user callback receives.

    Fields are ``None`` where the call site has nothing to offer:

    - ``x``: the current FLATTENED iterate (the whole parameter vector, not
      just the metric block).
    - ``lm_state``: the live :class:`LMState` (damping, ridge, caches, and
      the carried metric/preconditioner instances). In the implicit-AD rule
      this is the returned state under ``stop_gradient`` -- inert
      conditioning data, like the ridge.
    - ``args`` / ``p``: the residual's auxiliary data and differentiation
      parameters as passed to ``solve``/``update``.
    """

    x: Any = None
    lm_state: Any = None
    args: Any = None
    p: Any = None

nlls_gram.register_pytree_dataclass(cls, *, data_fields, meta_fields=())

Register a frozen dataclass as a pytree whose unflatten bypasses __init__/__post_init__.

data_fields become traced leaves (subtrees, if a field holds a container); meta_fields become static structure and must hold hashable values. Unflatten rebuilds the instance with object.__new__ plus object.__setattr__, so constructors are free to compute derived leaves and validate eagerly -- reconstruction inside jit, vmap, or a loop carry restores the stored fields verbatim without re-running any of it. Static values are type-tagged, so treedefs compare and hash by value with jit's strict-type semantics (1, 1.0, and True stay distinct). Every dataclass field must appear in exactly one of the two lists. Returns cls.

Every concrete :class:~nlls_gram.Metric and :class:~nlls_gram.Preconditioner class must be registered this way (subclassing a registered base does not register the subclass); the solvers reject unregistered instances at construction.

Source code in src/nlls_gram/utilities.py

def register_pytree_dataclass(cls, *, data_fields, meta_fields=()):
    """Register a frozen dataclass as a pytree whose unflatten bypasses
    ``__init__``/``__post_init__``.

    ``data_fields`` become traced leaves (subtrees, if a field holds a
    container); ``meta_fields`` become static structure and must hold
    hashable values. Unflatten rebuilds the instance with ``object.__new__``
    plus ``object.__setattr__``, so constructors are free to compute derived
    leaves and validate eagerly -- reconstruction inside jit, ``vmap``, or a
    loop carry restores the stored fields verbatim without re-running any of
    it. Static values are type-tagged, so treedefs compare and hash by value
    with jit's strict-type semantics (``1``, ``1.0``, and ``True`` stay
    distinct). Every dataclass field must appear in exactly one of the two
    lists. Returns ``cls``.

    Every concrete :class:`~nlls_gram.Metric` and
    :class:`~nlls_gram.Preconditioner` class must be registered this way
    (subclassing a registered base does not register the subclass); the
    solvers reject unregistered instances at construction.
    """
    data_fields = tuple(data_fields)
    meta_fields = tuple(meta_fields)
    declared = {f.name for f in dataclasses.fields(cls)}
    listed = set(data_fields) | set(meta_fields)
    if len(data_fields) + len(meta_fields) != len(listed) or listed != declared:
        raise ValueError(
            f"register_pytree_dataclass({cls.__name__}): data_fields + "
            f"meta_fields must cover every dataclass field exactly once; "
            f"declared {sorted(declared)}, listed {sorted(listed)}"
        )

    def static_aux(instance):
        # The typed tag makes equality/hash strict-typed; the raw value rides
        # alongside so unflatten can restore it verbatim.
        return tuple(
            (_typed_key(value), value)
            for value in (getattr(instance, name) for name in meta_fields)
        )

    def flatten_with_keys(instance):
        children = [
            (jax.tree_util.GetAttrKey(name), getattr(instance, name))
            for name in data_fields
        ]
        return children, static_aux(instance)

    def flatten(instance):
        return [getattr(instance, name) for name in data_fields], static_aux(instance)

    def unflatten(aux, children):
        instance = object.__new__(cls)
        for name, value in zip(data_fields, children, strict=True):
            object.__setattr__(instance, name, value)
        for name, (_, value) in zip(meta_fields, aux, strict=True):
            object.__setattr__(instance, name, value)
        return instance

    jax.tree_util.register_pytree_with_keys(cls, flatten_with_keys, unflatten, flatten)
    return cls

Multi-start

nlls_gram.MultiStart dataclass

Multi-start configuration for solve(multi_start=...).

draw(key, x, args) -> (x_new, args_new) generates a fresh initial condition; it must be traceable and type-stable (returning the same pytree structure, shapes, and dtypes as its (x, args) inputs). accept(key, result) -> bool optionally overrides the success test (default: CONVERGED plus MAX_STEPS when the solve's max_steps_is_success=True); it receives its own key so it can draw fresh validation data, and may return any scalar boolean-like value. Sequential mode (parallel=False) solves from (x0, args) and retries on failure, chaining each attempt's initial values into the next draw; parallel mode solves all num_starts lanes under vmap (lane 0 = the caller's (x0, args), the rest drawn from the originals) and selects the accepted lane with the lowest loss. The key schedule is draw_key, accept_key = jax.random.split(jax.random.fold_in(key, k)) for attempt k.

draw and accept enter the jit cache by identity (like callback): define them once at setup scope, not inline per call. MultiStart is not a pytree -- solve unpacks it before tracing, with key the only traced field.

Source code in src/nlls_gram/multi_start.py

@dataclass(frozen=True, eq=False)
class MultiStart:
    """Multi-start configuration for ``solve(multi_start=...)``.

    ``draw(key, x, args) -> (x_new, args_new)`` generates a fresh initial
    condition; it must be traceable and type-stable (returning the same pytree
    structure, shapes, and dtypes as its ``(x, args)`` inputs). ``accept(key,
    result) -> bool`` optionally overrides the success test (default:
    ``CONVERGED`` plus ``MAX_STEPS`` when the solve's
    ``max_steps_is_success=True``); it receives its own key so it can draw fresh
    validation data, and may return any scalar boolean-like value.
    Sequential mode (``parallel=False``) solves from ``(x0, args)`` and retries
    on failure, chaining each attempt's *initial* values into the next
    ``draw``; parallel mode solves all ``num_starts`` lanes under ``vmap``
    (lane 0 = the caller's ``(x0, args)``, the rest drawn from the originals)
    and selects the accepted lane with the lowest loss. The key schedule is
    ``draw_key, accept_key = jax.random.split(jax.random.fold_in(key, k))``
    for attempt ``k``.

    ``draw`` and ``accept`` enter the jit cache by identity (like
    ``callback``): define them once at setup scope, not inline per call.
    ``MultiStart`` is not a pytree -- ``solve`` unpacks it before tracing, with
    ``key`` the only traced field.
    """

    key: Any
    num_starts: int
    draw: Any = None
    accept: Any = None
    parallel: bool = False

    def __post_init__(self):
        if isinstance(self.num_starts, bool) or not isinstance(self.num_starts, int):
            raise ValueError("num_starts must be a Python int >= 1")
        if self.num_starts < 1:
            raise ValueError("num_starts must be a Python int >= 1")
        if self.num_starts > 1 and self.draw is None:
            raise ValueError(
                "num_starts > 1 requires draw; pass "
                "draw=(key, x, args) -> (x_new, args_new)"
            )
        if self.draw is not None and not callable(self.draw):
            raise TypeError("draw must be callable")
        if self.accept is not None and not callable(self.accept):
            raise TypeError("accept must be callable")

nlls_gram.MultiStartInfo dataclass

Diagnostics attached to LMSolveResult.multi_start by a multi-start solve.

attempt is the winning attempt/lane index (0 = the caller's (x0, args)), accepted whether the winner passed the success test (MultiStart.accept, or the solve's max_steps_is_success policy), and attempts_run how many starts were solved (sequential mode stops at the first success; parallel mode always runs num_starts). loss is the ranking objective used for selection -- the sum of squared residuals at the returned solution for LevenbergMarquardt, the ridge objective for RidgeLevenbergMarquardt -- masked to +inf when nonfinite. Note accepted describes the multi-start success test, not LMInfo.accepted (last-step acceptance).

Source code in src/nlls_gram/multi_start.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class MultiStartInfo:
    """Diagnostics attached to ``LMSolveResult.multi_start`` by a multi-start solve.

    ``attempt`` is the winning attempt/lane index (0 = the caller's
    ``(x0, args)``), ``accepted`` whether the winner passed the success test
    (``MultiStart.accept``, or the solve's ``max_steps_is_success`` policy), and
    ``attempts_run`` how many starts were solved (sequential mode stops at the
    first success; parallel mode always runs ``num_starts``). ``loss`` is the
    ranking objective used for selection -- the sum of squared residuals at the
    returned solution for ``LevenbergMarquardt``, the ridge objective for
    ``RidgeLevenbergMarquardt`` -- masked to ``+inf`` when nonfinite. Note
    ``accepted`` describes the multi-start success test, not ``LMInfo.accepted``
    (last-step acceptance).
    """

    attempt: jax.Array
    accepted: jax.Array
    attempts_run: jax.Array
    loss: jax.Array

nlls_gram.DrawNNXModule

Multi-start draw hook re-initializing a flax nnx.Module from a fresh key.

Given a MultiStart retry key, builds module_cls(*args, rngs=nnx.Rngs(key), **kwargs) and returns its nnx.Param state as the new solver start, passing args through unchanged. Use it instead of hand-rolling a re-init closure per driver::

draw = DrawNNXModule(SequentialMLP, settings, dtype=dtype)
ms = MultiStart(key=key, num_starts=5, draw=draw)

The drawn parameter state must be type-stable against the solver's x0 (same pytree structure, shapes, and dtypes) -- construct the module with a matching param_dtype/dtype (e.g. pass dtype= through). The paired nnx.GraphDef used by the residual's nnx.merge must come from the same module_cls(*args, **kwargs) spec.

Value-hashable on (module_cls, args, kwargs) with jit's strict-type semantics (1, 1.0, and True key distinct compilations): equal specs compare equal and share one jit compilation instead of recompiling per instance (a fresh closure would not). args/kwargs must be hashable for that sharing, and their values must not be mutated after construction (a stale key would reuse the wrong compile); unhashable specs still work but recompile per instance. Requires flax installed (imported lazily on first draw).

Source code in src/nlls_gram/multi_start.py

class DrawNNXModule:
    """Multi-start ``draw`` hook re-initializing a flax ``nnx.Module`` from a fresh key.

    Given a ``MultiStart`` retry key, builds
    ``module_cls(*args, rngs=nnx.Rngs(key), **kwargs)`` and returns its ``nnx.Param``
    state as the new solver start, passing ``args`` through unchanged. Use it instead
    of hand-rolling a re-init closure per driver::

        draw = DrawNNXModule(SequentialMLP, settings, dtype=dtype)
        ms = MultiStart(key=key, num_starts=5, draw=draw)

    The drawn parameter state must be type-stable against the solver's ``x0`` (same
    pytree structure, shapes, and dtypes) -- construct the module with a matching
    ``param_dtype``/``dtype`` (e.g. pass ``dtype=`` through). The paired
    ``nnx.GraphDef`` used by the residual's ``nnx.merge`` must come from the same
    ``module_cls(*args, **kwargs)`` spec.

    Value-hashable on ``(module_cls, args, kwargs)`` with jit's strict-type semantics
    (``1``, ``1.0``, and ``True`` key distinct compilations): equal specs compare equal
    and share one jit compilation instead of recompiling per instance (a fresh closure
    would not). ``args``/``kwargs`` must be hashable for that sharing, and their values
    must not be mutated after construction (a stale key would reuse the wrong compile);
    unhashable specs still work but recompile per instance. Requires ``flax`` installed
    (imported lazily on first draw).
    """

    def __init__(self, module_cls, *args, **kwargs):
        self.module_cls = module_cls
        self.args = args
        self.kwargs = tuple(sorted(kwargs.items()))

    def __call__(self, key, x_old, args_old):
        from flax import nnx

        module = self.module_cls(*self.args, rngs=nnx.Rngs(key), **dict(self.kwargs))
        _, theta = nnx.split(module, nnx.Param)
        return theta, args_old

    def __hash__(self):
        return hash((self.module_cls, _typed_key(self.args), _typed_key(self.kwargs)))

    def __eq__(self, other):
        return (
            isinstance(other, DrawNNXModule)
            and self.module_cls is other.module_cls
            and _typed_key(self.args) == _typed_key(other.args)
            and _typed_key(self.kwargs) == _typed_key(other.kwargs)
        )