Square Systems (Root Finding)

SquareLevenbergMarquardt is a solve-only damped-Newton (Levenberg-Marquardt) root solver for square nonsingular systems

\[ r(x, \text{args}, p) = 0, \qquad \dim r = \dim x , \]

built for hot inner loops — the motivating case is the algebraic stage solve of semi-explicit index-1 DAEs, where a small square root must be found at every Runge–Kutta stage, warm-started from the previous stage, and differentiated implicitly with respect to everything that defines the equation.

It deliberately does not reuse the underdetermined solver's machinery: at \(m = n\) the residual-space Gram \(J P J^\top\) squares the Jacobian's condition number for no benefit, and the general solver's metric, callback, and geodesic plumbing costs real time per stage. The square solver is a lean jitted lax.while_loop around a direct dense factorization.

API

from nlls_gram import SquareLevenbergMarquardt

solver = SquareLevenbergMarquardt(
    residual_fn,           # (x) | (x, args) | (x, args, p)
    init_damping=1e-3,
    damping_decrease=0.5,
    damping_increase=4.0,
    has_aux=False,
)

result = solver.solve(
    x0,              # root-finding guess only — never an AD target
    args,            # fixed data — never an AD target
    p=p,             # ALL differentiated equation inputs
    max_steps=8,
    atol=None,       # None -> 1e-6 (float32) / 1e-10 (float64)
)
z = result.x         # result.residual_norm, result.steps, result.status

There is no init/update pair — solve is the entire interface. The residual convention, x pytree flattening, and LMStatus codes match the underdetermined solver: x may be any pytree, and the residual must return a single array (or (array, aux) with has_aux=True, where aux is a pytree whose leaves are JAX numeric types — it is returned through the jitted loop) whose flattened size equals the flattened size of x, in the same dtype. The result is a lean SquareSolveResult with x, residual_norm (the residual 2-norm at the returned x), steps, status, aux (with has_aux=True), and the pass-through args and p the solve was called with — the same accessors as the underdetermined solver's LMSolveResult, so downstream code can read result.args/result.p uniformly. The loop controls (max_steps, atol, gtol, xtol) are concrete Python scalars, not traceable arguments.

For a DAE stage 0 = g(y, z, t, args, params) solved for z:

def residual(z, args, p):
    y, t, params = p
    return g(y, z, t, args, params)

result = solver.solve(z_guess, args, p=(y, t, params), max_steps=8, atol=root_atol)
z = result.x

The Step: Augmented QR, Never the Gram

Each iteration solves the LM subproblem \(\min_s \|r + J s\|_2^2 + \lambda \|s\|_2^2\) through one reduced QR of the augmented matrix

\[ \begin{bmatrix} J \\ \sqrt{\lambda}\, I \end{bmatrix} s \;\approx\; \begin{bmatrix} -r \\ 0 \end{bmatrix}, \]

whose normal equations are exactly \((J^\top J + \lambda I)s = -J^\top r\) — solved without ever forming \(J^\top J\) or \(J J^\top\), so the step does not square the Jacobian's condition number. For finite \(J\) and \(\lambda > 0\) the augmented matrix has full column rank (exact arithmetic), so the factorization does not fail: a singular Jacobian still yields a well-defined damped step, which the accept/reject test then vets like any other. Accept/reject on the sum of squared residuals and the multiplicative damping update follow the main solver.

The Jacobian is a dense forward-mode jax.jacfwd and is recomputed only after an accepted step (lax.cond); rejected steps reuse it. A warm start that already meets atol exits before the loop without computing a Jacobian at all — the DAE stage pattern (each stage starts near the next root) pays one residual evaluation for an already-converged guess (plus one more at the returned x when has_aux=True).

The linear algebra is deliberately dense and direct in this first version; support for other linear solvers may be added later behind a constructor knob.

Status Semantics

• LMStatus.CONVERGED with the default gtol=xtol=0 means exactly the residual criterion \(\|r\|_2 < \text{atol}\) — a DAE caller can trust the status as "found a root," and result.residual_norm is there to enforce a caller-side criterion directly.
• gtol (on \(\|J^\top r\|\)) and xtol (on an accepted step norm) are opt-in: they can report CONVERGED at a stagnation point that is not a root (e.g. the loss-stationary point of \(x^2 + 1\)), which is precisely why they default to off.
• NONFINITE reports a nonfinite residual at x0; a nonfinite candidate step is simply rejected (acceptance requires a finite loss). Exhausted budgets report MAX_STEPS. There are no runtime host exceptions under jit; the square-shape and dtype checks raise at trace time.

Implicit Differentiation

solve(...).x carries a custom implicit JVP with respect to p (the same jax.custom_jvp organization as the underdetermined solver): differentiate the defining equation at the returned root rather than the iterations. For \(r(x^*, \text{args}, p) = 0\),

\[ J_x \dot x = -J_p \dot p \]

is solved directly with the square dense \(J_x\) (LU) at the solution — never through \(J_x J_x^\top\). VJPs come from transposing the JVP rule, so forward, reverse, forward-over-reverse, and reverse-over-forward all compose (jax.hessian, jax.vmap over differentiated solves, etc.). For the DAE stage above this gives, with no extra rules in the caller,

\[ \dot z = -g_z^{-1}\left(g_y \dot y + g_t \dot t + g_\theta \dot\theta\right). \]

Contract notes:

• x0 and args receive zero tangents: the guess and fixed data are not AD targets. Anything that must be differentiated goes in p.
• Implicit derivatives are meaningful only when the returned point is a converged, nonsingular root — check status and residual_norm before trusting them. A singular \(J_x\) produces a non-finite tangent (loud), consistent with the library's conventions.
• result.residual_norm gets a zero tangent by contract (at a root it is ~0 and not meaningfully differentiable); with has_aux=True, result.aux is differentiated through both its direct p dependence and the root \(x^*(p)\).

Performance

Construct the residual once at setup scope — solvers compare by configuration, so a fresh equal-settings instance around the same residual function reuses the compiled loop, but a residual closure rebuilt per call is a new configuration and retraces, exactly as with the underdetermined solver. Measured on CPU (float32, 32 warm-started solves inside one jitted lax.scan, the DAE stage pattern; benchmarks/test_square_lm_benchmark.py), the full adaptive LM solve costs 1.0–1.6x a fixed-iteration direct-Newton fori_loop baseline at \(n \in \{1, 4, 8\}\) — within the issue #14 gate of ~2x, while adding damping robustness, early exit on atol, and the implicit-AD boundary.

API Reference

nlls_gram.SquareLevenbergMarquardt

Damped-Newton (Levenberg-Marquardt) root solver for square systems r(x, args, p) = 0 where the flattened residual and parameter sizes are equal and the Jacobian is nonsingular at the root.

The step solves the LM subproblem min_s ||r + J s||^2 + damping ||s||^2 through one reduced QR of the augmented matrix [J; sqrt(damping) I] -- a direct dense factorization that never forms J J' or J'J (no condition-number squaring), and whose augmented matrix has full column rank for any finite J when damping > 0 -- a singular Jacobian still yields a well-defined damped step, vetted by the accept/reject test, rather than a host exception. Only solve is exposed (no init/update); the loop is a lean lax.while_loop that evaluates the residual once per iteration plus once at x0 (and once more at the returned x when has_aux=True), recomputes the Jacobian only after an accepted step, and exits immediately on a warm start that already meets atol -- without ever computing a Jacobian. x may be any pytree; the residual must return a single array (or (array, aux) with has_aux=True, where aux is a pytree whose leaves are JAX numeric types) whose flattened size equals the flattened size of x, in the same dtype. Construct the solver once at setup scope; a new instance per call retraces.

solve(...).x has a custom implicit JVP with respect to p: the tangent solves the square system J_x xdot = -J_p pdot with a direct dense solve at the returned solution, and VJPs are obtained by transposition (jax.custom_jvp), so forward, reverse, and second-order differentiation all compose. x0 and args receive zero tangents by contract. Implicit derivatives are meaningful only when the returned point is a converged, nonsingular root -- check status and residual_norm.

Source code in src/nlls_gram/square_lm.py

class SquareLevenbergMarquardt:
    """Damped-Newton (Levenberg-Marquardt) root solver for square systems
    ``r(x, args, p) = 0`` where the flattened residual and parameter sizes
    are equal and the Jacobian is nonsingular at the root.

    The step solves the LM subproblem
    ``min_s ||r + J s||^2 + damping ||s||^2`` through one reduced QR of the
    augmented matrix ``[J; sqrt(damping) I]`` -- a direct dense factorization
    that never forms ``J J'`` or ``J'J`` (no condition-number squaring), and
    whose augmented matrix has full column rank for any finite ``J`` when
    ``damping > 0`` -- a singular Jacobian still yields a well-defined damped
    step, vetted by the accept/reject test, rather than a host exception.
    Only ``solve`` is exposed (no ``init``/``update``); the
    loop is a lean ``lax.while_loop`` that evaluates the residual once per
    iteration plus once at ``x0`` (and once more at the returned ``x`` when
    ``has_aux=True``), recomputes the Jacobian only after an accepted step,
    and exits immediately on a warm start that already meets ``atol`` --
    without ever computing a Jacobian. ``x`` may be any pytree; the residual
    must return a single array (or ``(array, aux)`` with ``has_aux=True``,
    where ``aux`` is a pytree whose leaves are JAX numeric types) whose
    flattened size equals the flattened size of ``x``, in the same
    dtype. Construct the solver once at setup scope; a new instance per call
    retraces.

    ``solve(...).x`` has a custom implicit JVP with respect to ``p``: the
    tangent solves the square system ``J_x xdot = -J_p pdot`` with a direct
    dense solve at the returned solution, and VJPs are obtained by
    transposition (``jax.custom_jvp``), so forward, reverse, and second-order
    differentiation all compose. ``x0`` and ``args`` receive zero tangents by
    contract. Implicit derivatives are meaningful only when the returned
    point is a converged, nonsingular root -- check ``status`` and
    ``residual_norm``.
    """

    def __init__(
        self,
        residual_fn,
        *,
        init_damping=1e-3,
        damping_decrease=0.5,
        damping_increase=4.0,
        has_aux=False,
    ):
        canonical_residual, residual_arity = canonicalize_residual(residual_fn)
        if init_damping <= 0:
            raise ValueError("init_damping must be positive")
        if damping_decrease <= 0:
            raise ValueError("damping_decrease must be positive")
        if damping_increase <= 0:
            raise ValueError("damping_increase must be positive")
        self.residual_fn = canonical_residual
        self.residual_arity = residual_arity
        self.init_damping = init_damping
        self.damping_decrease = damping_decrease
        self.damping_increase = damping_increase
        self.has_aux = has_aux
        # Value-based identity: the jitted solve loop marks the solver itself
        # static, so equal-config solvers built around the same residual share
        # the compiled loop across instances instead of retracing per
        # construction.
        self._static_key = tuple(
            _static_key_component(value)
            for value in (
                residual_fn,
                init_damping,
                damping_decrease,
                damping_increase,
                has_aux,
            )
        )
        self._static_hash = hash(self._static_key)

    def __eq__(self, other):
        if self is other:
            return True
        if type(other) is not type(self):
            return NotImplemented
        return self._static_key == other._static_key

    def __hash__(self):
        return self._static_hash

    def solve(
        self,
        x0,
        args=None,
        *,
        p=None,
        max_steps=256,
        atol=None,
        gtol=0.0,
        xtol=0.0,
        jit=True,
    ):
        """Run damped-Newton steps from ``x0`` until a stopping rule fires.

        ``atol`` stops on the residual 2-norm; ``None`` uses a dtype-aware
        default (``1e-6`` in float32, ``1e-10`` in float64). ``gtol`` stops on
        the gradient norm ``||J' r||`` and ``xtol`` on an accepted step norm;
        both default to ``0`` (disabled), so ``LMStatus.CONVERGED`` means the
        residual criterion unless they are opted into. ``x0`` is a
        root-finding guess only; a warm start already at the root returns in
        zero steps. The loop controls (``max_steps``, ``atol``, ``gtol``,
        ``xtol``) are concrete Python scalars validated eagerly -- they are
        not traceable arguments.
        """
        # Silently dropping args/p a residual never sees would, in particular,
        # make the implicit derivative with respect to p a silent zero.
        if args is not None and self.residual_arity < 2:
            raise ValueError(
                "args was passed but residual_fn takes only (x); "
                "use residual_fn(x, args)"
            )
        if p is not None and self.residual_arity < 3:
            raise ValueError(
                "p was passed but residual_fn takes no p argument; "
                "use residual_fn(x, args, p)"
            )
        if not isinstance(max_steps, int) or isinstance(max_steps, bool):
            raise ValueError("max_steps must be a positive int")
        if max_steps <= 0:
            raise ValueError("max_steps must be a positive int")
        if atol is not None and atol < 0:
            raise ValueError("atol must be nonnegative or None")
        if gtol < 0:
            raise ValueError("gtol must be nonnegative")
        if xtol < 0:
            raise ValueError("xtol must be nonnegative")

        @jax.custom_jvp
        def solve_with_implicit_p(x, args, p, max_steps, atol, gtol, xtol):
            if jit:
                return _square_solve_loop_jit(
                    self, x, args, p, max_steps, atol, gtol, xtol
                )
            return self._solve_python(x, args, p, max_steps, atol, gtol, xtol)

        @solve_with_implicit_p.defjvp
        def solve_with_implicit_p_jvp(primals, tangents):
            x, args, p, max_steps, atol, gtol, xtol = primals
            _, _, p_dot, _, _, _, _ = tangents
            result = solve_with_implicit_p(x, args, p, max_steps, atol, gtol, xtol)
            x_dot = self._implicit_x_tangent_from_p(result.x, args, p, p_dot)
            zero_result = jax.tree.map(_zero_tangent_leaf, result)
            aux_dot = zero_result.aux
            if self.has_aux and p is not None:
                # aux depends on p directly and through the root x*(p);
                # linearize the aux map at the returned solution to account
                # for both paths.
                def aux_at_solution(x_value, p_value):
                    return self.residual_fn(x_value, args, p_value)[1]

                aux_dot = jax.jvp(aux_at_solution, (result.x, p), (x_dot, p_dot))[1]
            return (
                result,
                SquareSolveResult(
                    x_dot,
                    zero_result.residual_norm,
                    zero_result.steps,
                    zero_result.status,
                    aux_dot,
                    zero_result.args,
                    p_dot,
                ),
            )

        return solve_with_implicit_p(x0, args, p, max_steps, atol, gtol, xtol)

    def _implicit_x_tangent_from_p(self, x, args, p, p_dot):
        if p is None:
            return jax.tree.map(_zero_tangent_leaf, x)
        theta, unravel = ravel_pytree(x)

        def residual_from_theta(theta_value):
            value = self.residual_fn(unravel(theta_value), args, p)
            if self.has_aux:
                value = value[0]
            return jnp.ravel(value)

        def residual_from_p(p_value):
            value = self.residual_fn(x, args, p_value)
            if self.has_aux:
                value = value[0]
            return jnp.ravel(value)

        J_x = jax.jacfwd(residual_from_theta)(theta)
        residual_p_dot = jax.jvp(residual_from_p, (p,), (p_dot,))[1]
        theta_dot = -jnp.linalg.solve(J_x, residual_p_dot)
        return unravel(theta_dot)

    def _solve_python(self, x, args, p, max_steps, atol, gtol, xtol):
        setup = _square_solve_setup(self, x, args, p, atol, gtol, xtol)
        theta, unravel, residual_flat, resid, tolerances = setup
        atol, gtol, xtol = tolerances
        n = theta.size
        dtype = resid.dtype
        damping = jnp.asarray(self.init_damping, dtype=dtype)
        loss = jnp.sum(resid**2)
        status = LMStatus.RUNNING
        steps = 0
        J = None
        if not bool(jnp.isfinite(loss)):
            status = LMStatus.NONFINITE
        elif bool((atol > 0) & (jnp.sqrt(loss) < atol)):
            status = LMStatus.CONVERGED

        while status == LMStatus.RUNNING and steps < max_steps:
            if J is None:
                J = jax.jacfwd(residual_flat)(theta)
            grad_norm = jnp.linalg.norm(J.T @ resid)
            s = _square_damped_step(J, resid, damping, n, dtype)
            resid_candidate = residual_flat(theta + s)
            loss_candidate = jnp.sum(resid_candidate**2)
            improved = bool(jnp.isfinite(loss_candidate) & (loss_candidate < loss))
            if improved:
                theta = theta + s
                resid = resid_candidate
                loss = loss_candidate
                J = None
                damping = damping * self.damping_decrease
            else:
                damping = damping * self.damping_increase
            steps += 1
            if not bool(jnp.isfinite(loss)):
                status = LMStatus.NONFINITE
            elif bool(
                ((atol > 0) & (jnp.sqrt(loss) < atol))
                | ((gtol > 0) & (grad_norm < gtol))
                | ((xtol > 0) & improved & (jnp.linalg.norm(s) < xtol))
            ):
                status = LMStatus.CONVERGED

        if status == LMStatus.RUNNING:
            status = LMStatus.MAX_STEPS
        aux = None
        if self.has_aux:
            aux = self.residual_fn(unravel(theta), args, p)[1]
        return SquareSolveResult(
            unravel(theta),
            jnp.sqrt(jnp.sum(resid**2)),
            jnp.asarray(steps, dtype=jnp.int32),
            jnp.asarray(status, dtype=jnp.int32),
            aux,
            args,
            p,
        )

solve(x0, args=None, *, p=None, max_steps=256, atol=None, gtol=0.0, xtol=0.0, jit=True)

Run damped-Newton steps from x0 until a stopping rule fires.

atol stops on the residual 2-norm; None uses a dtype-aware default (1e-6 in float32, 1e-10 in float64). gtol stops on the gradient norm ||J' r|| and xtol on an accepted step norm; both default to 0 (disabled), so LMStatus.CONVERGED means the residual criterion unless they are opted into. x0 is a root-finding guess only; a warm start already at the root returns in zero steps. The loop controls (max_steps, atol, gtol, xtol) are concrete Python scalars validated eagerly -- they are not traceable arguments.

Source code in src/nlls_gram/square_lm.py

def solve(
    self,
    x0,
    args=None,
    *,
    p=None,
    max_steps=256,
    atol=None,
    gtol=0.0,
    xtol=0.0,
    jit=True,
):
    """Run damped-Newton steps from ``x0`` until a stopping rule fires.

    ``atol`` stops on the residual 2-norm; ``None`` uses a dtype-aware
    default (``1e-6`` in float32, ``1e-10`` in float64). ``gtol`` stops on
    the gradient norm ``||J' r||`` and ``xtol`` on an accepted step norm;
    both default to ``0`` (disabled), so ``LMStatus.CONVERGED`` means the
    residual criterion unless they are opted into. ``x0`` is a
    root-finding guess only; a warm start already at the root returns in
    zero steps. The loop controls (``max_steps``, ``atol``, ``gtol``,
    ``xtol``) are concrete Python scalars validated eagerly -- they are
    not traceable arguments.
    """
    # Silently dropping args/p a residual never sees would, in particular,
    # make the implicit derivative with respect to p a silent zero.
    if args is not None and self.residual_arity < 2:
        raise ValueError(
            "args was passed but residual_fn takes only (x); "
            "use residual_fn(x, args)"
        )
    if p is not None and self.residual_arity < 3:
        raise ValueError(
            "p was passed but residual_fn takes no p argument; "
            "use residual_fn(x, args, p)"
        )
    if not isinstance(max_steps, int) or isinstance(max_steps, bool):
        raise ValueError("max_steps must be a positive int")
    if max_steps <= 0:
        raise ValueError("max_steps must be a positive int")
    if atol is not None and atol < 0:
        raise ValueError("atol must be nonnegative or None")
    if gtol < 0:
        raise ValueError("gtol must be nonnegative")
    if xtol < 0:
        raise ValueError("xtol must be nonnegative")

    @jax.custom_jvp
    def solve_with_implicit_p(x, args, p, max_steps, atol, gtol, xtol):
        if jit:
            return _square_solve_loop_jit(
                self, x, args, p, max_steps, atol, gtol, xtol
            )
        return self._solve_python(x, args, p, max_steps, atol, gtol, xtol)

    @solve_with_implicit_p.defjvp
    def solve_with_implicit_p_jvp(primals, tangents):
        x, args, p, max_steps, atol, gtol, xtol = primals
        _, _, p_dot, _, _, _, _ = tangents
        result = solve_with_implicit_p(x, args, p, max_steps, atol, gtol, xtol)
        x_dot = self._implicit_x_tangent_from_p(result.x, args, p, p_dot)
        zero_result = jax.tree.map(_zero_tangent_leaf, result)
        aux_dot = zero_result.aux
        if self.has_aux and p is not None:
            # aux depends on p directly and through the root x*(p);
            # linearize the aux map at the returned solution to account
            # for both paths.
            def aux_at_solution(x_value, p_value):
                return self.residual_fn(x_value, args, p_value)[1]

            aux_dot = jax.jvp(aux_at_solution, (result.x, p), (x_dot, p_dot))[1]
        return (
            result,
            SquareSolveResult(
                x_dot,
                zero_result.residual_norm,
                zero_result.steps,
                zero_result.status,
                aux_dot,
                zero_result.args,
                p_dot,
            ),
        )

    return solve_with_implicit_p(x0, args, p, max_steps, atol, gtol, xtol)

nlls_gram.SquareSolveResult dataclass

Result of SquareLevenbergMarquardt.solve.

residual_norm is the residual 2-norm at the returned x -- with the default tolerances, status == LMStatus.CONVERGED holds exactly when it is below atol, and a caller enforcing its own root criterion can check it directly. aux is the residual's aux output at the returned x (has_aux=True only). args and p are the values the solve was called with, carried through for a uniform result interface with LMSolveResult (the square solver has no callbacks, so they are exactly the inputs; p participates in the implicit rule like the underdetermined solver's).

Source code in src/nlls_gram/square_lm.py

@jax.tree_util.register_dataclass
@dataclass(frozen=True)
class SquareSolveResult:
    """Result of ``SquareLevenbergMarquardt.solve``.

    ``residual_norm`` is the residual 2-norm at the returned ``x`` -- with the
    default tolerances, ``status == LMStatus.CONVERGED`` holds exactly when it
    is below ``atol``, and a caller enforcing its own root criterion can check
    it directly. ``aux`` is the residual's aux output at the returned ``x``
    (``has_aux=True`` only). ``args`` and ``p`` are the values the solve was
    called with, carried through for a uniform result interface with
    ``LMSolveResult`` (the square solver has no callbacks, so they are exactly
    the inputs; ``p`` participates in the implicit rule like the
    underdetermined solver's).
    """

    x: Any
    residual_norm: jax.Array
    steps: jax.Array
    status: jax.Array
    aux: Any = None
    args: Any = None
    p: Any = None