Skip to content

Metric and Preconditioner Utilities

The library ships a small set of constructors and helpers so that models can assemble metrics and CG preconditioners from structure they already know, instead of hand-rolling callback plumbing. Everything here returns plain callables or Metric objects that the solver sees through the metric=, dual_preconditioner=, and implicit_preconditioner= arguments. The CG paths require explicit preconditioners; identity_preconditioner() is the explicit opt-out.

Helper Builds Cost per apply
metric_from_cholesky(L) dense Metric from \(M = LL^\top\) \(O(n^2)\)
metric_from_tridiagonal_precision(diag, off_diag) Metric from a tridiagonal \(T = M^{-1}\) \(O(n)\)
metric_from_state_space(points, h, Pinf, transition) Metric for a stationary state-space kernel Gram \(M = K + \eta I\) (Matérn via matern_state_space) \(O(n m^2)\)
metric_from_quasiseparable(d, p, q, A) Metric from quasiseparable generators \(O(n m^2)\)
metric_from_shifted_matvec(matvec, shift) matrix-free Metric for \(M = A + \varepsilon I\) via inner CG \(O(\text{iters} \times \text{matvec})\)
metric_from_diagonal(weights) Metric from \(M = \operatorname{diag}(w)\) \(O(n)\)
blockdiag_metric(blocks) Metric over concatenated parameter blocks sum of blocks
sherman_morrison_preconditioner(solve, u, weight) dual_preconditioner for \(B = A + w\,uu^\top\) one solve
woodbury_preconditioner(solve, U, weights) dual_preconditioner for \(B = A + U\operatorname{diag}(w)U^\top\) one solve + \(k \times k\)
identity_preconditioner() the explicit "no preconditioner" choice (both hook signatures) free
nystrom_preconditioner(matvec, n, rank, key) randomized Nyström dual_preconditioner for a PSD operator (FTU) two \((n, \text{rank})\) GEMVs
pad_dual_preconditioner(base, n_real) extends a dual_preconditioner to a zero-padded residual base + \(O(k)\)

Tridiagonal Precision Metric

For Markov kernels the Gram inverse is exactly tridiagonal — for the Matérn-1/2 / Ornstein-Uhlenbeck kernel on sorted points \(t_1 < \dots < t_n\) with \(\rho_i = e^{-(t_{i+1}-t_i)/\ell}\), the precision has closed-form entries. Passing the two diagonals gives a Metric whose every callback is \(O(n)\), with nothing factored densely:

from nlls_gram import metric_from_tridiagonal_precision

metric = metric_from_tridiagonal_precision(diag, off_diag)

parallel=None (the default) runs the one-time bidiagonal Cholesky setup as an associative \(O(\log n)\)-depth scan off-CPU in float64 — where a sequential scan pays a kernel launch per step — and as the sequential scan otherwise. In float32 the default stays sequential even off-CPU: the parallel scan's projective \(2\times 2\) products can cancel to non-finite values on long, stiff grids (near-unit-correlation AR(1)), while the sequential recurrence is stable there.

State-Space Kernel Metrics (Quasiseparable)

A stationary Gaussian process has an exact O(n) Gram factorization precisely when it admits a finite-dimensional state-space (linear SDE) representation: an \(m\)-dimensional latent Gauss-Markov state observed through a row vector \(h\). With stationary state covariance \(P_\infty\) and transition matrices \(A_k = \Phi(t_k - t_{k-1})^\top\) (transposed matrix exponential of the SDE drift, the tinygp orientation), the Gram on sorted points is

\[ K_{ij} = h^\top P_\infty A_i A_{i-1} \cdots A_{j+1}\, h \quad (i > j), \]

\(h^\top P_\infty h\) on the diagonal, symmetric — a quasiseparable (celerite-style, rank-\(m\) semiseparable) matrix whose Cholesky factor shares the same structure, so every callback is one or two O(\(n m^2\)) scans.

The main application is the half-integer Matérn family, which is exactly the CAR(\(m\)) state-space class with \(m = 1, 2, 3\) for \(\nu = 1/2, 3/2, 5/2\); matern_state_space(sigma, ell, nu) supplies the exact \((h, P_\infty, \Phi^\top)\) mapping (with \(f = \sqrt{2\nu}/\ell\): \(h = [\sigma, 0, \ldots]\), e.g. \(P_\infty = \operatorname{diag}(1, f^2)\) for \(\nu = 3/2\)). This route is necessary for exactness: only the Matérn-1/2 value Gram has a sparse inverse. For 3/2 and 5/2 the sampled value process is ARMA-like, not Markov — the Gram-inverse off-band entries are \(\sim 10^{-2}\) relative (versus \(\sim 10^{-16}\) for 1/2), so a truncated band would be an approximate metric, and the library's contract requires metric.solve to be exact (approximations belong in dual_preconditioner). Only the latent state is Markov; the state-space form exploits exactly that.

from nlls_gram import matern_state_space, metric_from_state_space

metric = metric_from_state_space(
    points, *matern_state_space(sigma, ell, nu=1.5), nugget=1e-8 * sigma**2
)

Other stationary state-space kernels (sums of exponentials, CARMA / celerite-style terms) drop into the same constructor through their own \((h, P_\infty, \Phi^\top)\).

points must be 1-D and sorted strictly increasing — not validated, since it may be traced; unsorted or repeated points silently produce a wrong or NaN metric. The metric is \(M = K + \eta I\): the absolute nugget \(\eta\) folds into the diagonal generator before factorization, so it is part of the metric — exact, not a solver fudge. Nugget-free Matérn-3/2 and 5/2 Grams on fine grids are extremely ill-conditioned (condition number \(\sim 10^{21}\) at \(n = 5000\) — a property of the matrix, not the solver); supply a nugget whenever the grid resolves the kernel. For \(\nu = 1/2\) the constructor works but metric_from_tridiagonal_precision is the specialized alternative whose applies are elementwise shifts — strictly cheaper than scans on GPU.

metric_from_quasiseparable(d, p, q, A, nugget=0.0, parallel=None) is the generator-level general API — any stationary state-space kernel (sums of exponentials, celerite terms) reduces to it, and banded matrices are themselves rank-\(p\) quasiseparable. \(A_k\) is the transition INTO index \(k\) (\(A_0\) never enters the products; the state-space builders set it to the identity). Positive definiteness is not validated (inputs may be traced): a non-PD input silently produces NaN through the Cholesky square roots, the same convention as the tridiagonal constructor.

parallel=None (the default) picks the apply implementation once at construction, from the backend and dtype there: associative O(\(\log n\))-depth scans off-CPU in float64 — where a sequential scan pays a kernel launch per step — and sequential scans otherwise. Unlike the tridiagonal constructor's setup scan, the parallel substitutions here propagate rank-1-corrected transition matrices \(A_k - w_k p_k^\top / c_k\) with no contraction guarantee, so the float32 default stays sequential on every backend; the float64 default is backed by the stress-grid agreement tests. Pass parallel=True/False to force either path. The one-time Cholesky setup is a sequential scan in this release — cheap for fixed metrics reused across solves, but on the hot path when the metric is rebuilt from traced \(\sigma, \ell\) inside jax.grad/vmap sweeps; see the Tuning Guide.

Unified Shifted Block Metrics

For a kernel-coefficient block \(\alpha\) plus scalar parameters \(\beta\), the unified shifted metric \(M = \operatorname{blockdiag}(K, 0) + \varepsilon I = \operatorname{blockdiag}(K + \varepsilon I_n, \varepsilon I_k)\) completes the RKHS seminorm \(\alpha^\top K \alpha\) with a single spectral floor — the theory (spectrum, the \(\varepsilon \to 0\) seminorm limit and its uniqueness conditions) is in Metric Gauss-Newton. It composes from existing constructors; the kernel block has three interchangeable representations:

scalar_block = metric_from_diagonal(eps * jnp.ones(k))
metric = blockdiag_metric([(kernel_block, n), (scalar_block, k)])
# kernel_block, by representation of K:
#   metric_from_cholesky(jnp.linalg.cholesky(K + eps * jnp.eye(n)))   # dense
#   metric_from_state_space(t, *matern_state_space(sigma, ell, nu), nugget=eps)
#   metric_from_shifted_matvec(kernel_matvec, eps)      # matvec only; cholesky/cg

One structural note: the Matérn-1/2 tridiagonal shortcut does not survive the shift — metric_from_tridiagonal_precision parameterizes \(T = M^{-1}\), and \((K + \varepsilon I)^{-1}\) is not tridiagonal — so a shifted OU metric goes through metric_from_state_space(..., nu=0.5, nugget=eps) instead.

metric_from_shifted_matvec(matvec, shift, *, tol=None, atol=0.0, maxiter=None, preconditioner=None) needs only a matvec of a symmetric PSD \(A\), accepting (n,) and (n, k) leading-axis inputs (the same shape contract Metric.solve carries). The positive shift is what makes an iterative metric solve viable: \(\kappa(A + \varepsilon I) \le (\lambda_{\max} + \varepsilon)/\varepsilon\) regardless of how singular \(A\) is — and in practice far better, since the shift clusters the spectral tail at \(\approx \varepsilon\) and CG resolves a cluster in about one iteration (measured: ~32 float64 iterations for a Matérn-5/2 Gram at n=1000, independent of \(\varepsilon\) from 1e-2 to 1e-8). It provides solve and norm only (no matrix-free square root), so it works with the cholesky and cg linear solvers and is rejected for qr at construction. Combined with implicit_solver="cg" the whole pipeline — forward solve, JVP, and VJP — runs matrix-free (see Implicit AD). This is the one constructor that meets the metric.solve exactness contract in a limit rather than identically: its inner CG tolerance is part of the answer, not of the schedule — the residual error perturbs the selected solution and the implicit derivatives at order tol, and the implicit derivative has no accept/reject safeguard. The default tolerance (square root of the dtype's machine epsilon) matches typical outer tolerances; never cap maxiter as a cost control (a truncated CG is not a linear map, which breaks the cg linear_solver's operator assumptions).

Diagonal and Block-Diagonal Metrics

Models with a kernel block plus a few scalar parameters can compose per-block metrics instead of writing slice/concatenate glue. The blocks are laid out in the order the solver flattens the parameter pytree (ravel_pytree order):

import jax.numpy as jnp

from nlls_gram import blockdiag_metric, metric_from_cholesky, metric_from_diagonal

metric = blockdiag_metric(
    [
        (metric_from_cholesky(jnp.linalg.cholesky(K)), n),
        (metric_from_diagonal(jnp.full(1, m_0)), 1),
    ]
)

solve, inv_sqrt, and inv_sqrt_transpose slice on the leading axis, so vector and matrix inputs both work; norm combines the block norms in quadrature. A fully-default Metric() block means the identity metric on that block. A block that defines some callbacks but leaves others None propagates the missing callbacks as None on the composite, so the solver's construction-time validation applies exactly as it would to that block alone.

Sherman–Morrison Dual Preconditioner

With linear_solver="cg", the dual_preconditioner(v, damping) argument supplies an approximation of \((J M^{-1} J^\top + \lambda I)^{-1} v\) on residual-space vectors. It never changes the subproblem being solved: at inner convergence the step is identical, and a budget-truncated step still lies in \(\operatorname{range}(M^{-1}J^\top)\), so the minimum-metric-norm selection for underdetermined residuals is unchanged and an approximate preconditioner is safe — even though metric.solve must stay exact.

A metric weight \(m\) on a scalar parameter injects an exactly known rank-1 spike into the dual operator. For the kernel-collocation family, a Jacobian column \(-c\,u\) for that parameter contributes \((c^2/m)\,uu^\top\), and

from nlls_gram import sherman_morrison_preconditioner

dual_preconditioner = sherman_morrison_preconditioner(
    alpha_metric.solve, jnp.ones(n), c**2 / m_0
)

builds \(B^{-1}\) for \(B = K + (c^2/m_0)\,\mathbf{1}\mathbf{1}^\top\) from one kernel solve plus a rank-1 correction (\(B\), not \(P\): the docs reserve \(P\) for \(M^{-1}\)). Under the unified shifted metric the scalar-block weight is \(\varepsilon\), so the spike weight is \(c^2/\varepsilon\) — it grows as \(\varepsilon\) shrinks, making the preconditioner more load-bearing, not less. For \(k\) scalar parameters at once, woodbury_preconditioner(solve, U, weights) is the rank-\(k\) generalization (\(B = A + U\operatorname{diag}(w)U^\top\), one matrix solve plus a \(k \times k\) Cholesky, reducing exactly to Sherman-Morrison at \(k = 1\)). Such structural preconditioners can be spectrally equivalent to the dual operator uniformly in \(n\), keeping the inner CG budget constant where the unpreconditioned budget grows with refinement — see the Tuning Guide.

Identity Preconditioner

linear_solver="cg" requires a dual_preconditioner, and a cg-resolved implicit solve requires an implicit_preconditioner — running Krylov methods unpreconditioned should be a decision, not a default. identity_preconditioner() is that decision made explicit and greppable:

from nlls_gram import identity_preconditioner

solver = UnderdeterminedLevenbergMarquardt(
    residual_fn,
    linear_solver="cg",
    dual_preconditioner=identity_preconditioner(),
    implicit_preconditioner=identity_preconditioner(),
)

The returned callable accepts both hook signatures — dual_preconditioner(v, damping) and implicit_preconditioner(v) — so one helper serves both arguments.

Nyström Preconditioner for Neural-Network Least Squares

When no structural preconditioner is available — typically neural-network least squares under the identity metric, where the dual operator is the \(m \times m\) empirical NTK Gram \(JJ^\top\) with fast spectral decay — nystrom_preconditioner(matvec, n, rank, key) builds the randomized Nyström preconditioner of Frangella, Tropp, and Udell: sketch the PSD operator with a thin-QR'd Gaussian test matrix (rank operator applications plus one \(O(n\,\text{rank}^2)\) factorization, once at construction), recover the rank-rank approximation \(\hat A = U\Lambda U^\top\), and apply

\[ v \mapsto U\frac{U^\top v}{\Lambda + \lambda} + \frac{v - UU^\top v}{\rho + \lambda}, \]

with \(\rho\) the smallest retained eigenvalue and \(\lambda\) the live LM damping — this is the one shipped helper that uses the damping argument, so a single construction serves every damping value. The unresolved complement is balanced at \(\rho + \lambda\) rather than \(\lambda\); that balance carries the FTU condition-number guarantee for fast-decaying spectra. matvec must apply a symmetric PSD operator and accept (n, k) matrices (the Metric.solve shape contract). Like every preconditioner it is frozen at construction, so for a nonlinear residual it approximates the dual at the linearization point it was built from — staleness across LM steps is safe (preconditioner error never moves the converged root), and refresh cadence is a tuning knob.

The NTK matvec assembles matrix-free from the residual at the initial parameters via jax.linearize and jax.linear_transpose (this mirrors the test_cg_nystrom_mlp_ntk_example test, which doubles as the runnable example):

import jax
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree

from nlls_gram import (
    UnderdeterminedLevenbergMarquardt,
    identity_preconditioner,
    nystrom_preconditioner,
)

# residual: m collocation residuals of a pure-jax MLP, n_params >> m
theta0, unravel = ravel_pytree(x0)
_, jvp_fn = jax.linearize(lambda th: residual(unravel(th)), theta0)
transpose_fn = jax.linear_transpose(jvp_fn, theta0)


def ntk_matvec(V):  # (m, k) -> J (J' V), frozen at x0
    return jax.vmap(
        lambda col: jvp_fn(transpose_fn(col)[0]), in_axes=1, out_axes=1
    )(V)


solver = UnderdeterminedLevenbergMarquardt(
    residual,
    linear_solver="cg",
    iterative_tol=1e-6,
    iterative_maxiter=20,
    dual_preconditioner=nystrom_preconditioner(
        ntk_matvec, m, rank, jax.random.PRNGKey(0)
    ),
    implicit_preconditioner=identity_preconditioner(),
)

Passed as implicit_preconditioner the helper applies its undamped (zero-damping) inverse, which is valid only when the retained spectrum is strictly positive.

Padded Zero Residuals (Fixed Residual Shape)

Some JAX workflows keep a fixed residual shape across problem instances by appending residual entries that are identically zero:

def residual_padded(x):
    r = residual(x)
    return jnp.concatenate((r, jnp.zeros(pad, r.dtype)))

The padded rows have zero Jacobian rows, so the residual-space dual operator becomes exactly block diagonal — \(\operatorname{blockdiag}(J P J^\top + \lambda I,\; \lambda I)\) — and the solvers behave as follows:

  • cholesky is unchanged mathematically: the padded block decouples exactly, and the step matches the unpadded step (regression-tested for both the plain and geodesic-accelerated updates). The cost is the larger materialized residual dimension and dual factor; a few padded rows are harmless, large padding pays the dense residual-space price.

  • cg: a shape-fixed dual_preconditioner (a dense solve, nystrom_preconditioner, or a Sherman-Morrison/Woodbury built at the unpadded size) fails on the padded residual space and must be wrapped; pad_dual_preconditioner(base_preconditioner, n_real) applies the base callback on the first n_real coordinates and the exact \(1/\lambda\) inverse on the padded block:

    from nlls_gram import pad_dual_preconditioner

dual_preconditioner = pad_dual_preconditioner(base_preconditioner, n_real)

A shape-generic base (identity_preconditioner()) stays valid unwrapped — it just forgoes the exact padded-block inverse. Do not zero the padded block instead: that makes the preconditioner singular rather than SPD, even though it can appear to work when exactly-zero padding never excites those coordinates. - qr does not survive padding: the padded zero rows make the Jacobian rank-deficient, which the QR path's triangular solves cannot handle — the step is non-finite. Use cholesky or cg for padded problems. - Implicit AD: the padded rows make the undamped implicit dual \(J P J^\top\) singular, so the library's implicit rules (dense and cg) return a non-finite derivative of solve(...).x on padded problems — and for the same reason pad_dual_preconditioner divides the padded block by the live damping and is rejected at construction when passed as an implicit_preconditioner. The minimum-metric-norm derivative still exists mathematically (padding only appends redundant equations) and equals the unpadded derivative, so differentiate the unpadded formulation to compute it.

API

nlls_gram.metric_from_tridiagonal_precision(diag, off_diag, parallel=None)

Build an O(n) Metric from a symmetric tridiagonal precision matrix.

diag (length n) and off_diag (length n-1) give the main and off-diagonal of T = M^{-1}, which must be positive definite (not validated, since inputs may be traced; non-PD input silently produces a non-metric). Suited to Markov kernels, where the Gram inverse is exactly tridiagonal (e.g. Matern-1/2 / Ornstein-Uhlenbeck on sorted points): every callback costs O(n) and nothing is factored densely.

parallel picks how the one-time bidiagonal Cholesky setup runs: None (default) uses an associative O(log n)-depth scan off-CPU in float64 — where a sequential scan pays a kernel launch per step — and the sequential scan otherwise. In float32 the default stays sequential even off-CPU: the parallel scan evaluates the pivot recurrence through projective 2x2 products whose cancellation can go non-finite on long, stiff grids (e.g. near-unit-correlation AR(1)), while the sequential recurrence is stable there. Pass parallel=True to override only when the setup is float64 or the grid is short and well-conditioned.

Source code in src/nlls_gram/metrics.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def metric_from_tridiagonal_precision(diag, off_diag, parallel=None):
    """Build an O(n) ``Metric`` from a symmetric tridiagonal precision matrix.

    ``diag`` (length n) and ``off_diag`` (length n-1) give the main and
    off-diagonal of ``T = M^{-1}``, which must be positive definite (not
    validated, since inputs may be traced; non-PD input silently produces a
    non-metric). Suited to
    Markov kernels, where the Gram inverse is exactly tridiagonal (e.g.
    Matern-1/2 / Ornstein-Uhlenbeck on sorted points): every callback costs
    O(n) and nothing is factored densely.

    ``parallel`` picks how the one-time bidiagonal Cholesky setup runs:
    ``None`` (default) uses an associative O(log n)-depth scan off-CPU in
    float64 — where a sequential scan pays a kernel launch per step — and
    the sequential scan otherwise. In float32 the default stays sequential
    even off-CPU: the parallel scan evaluates the pivot recurrence through
    projective 2x2 products whose cancellation can go non-finite on long,
    stiff grids (e.g. near-unit-correlation AR(1)), while the sequential
    recurrence is stable there. Pass ``parallel=True`` to override only when
    the setup is float64 or the grid is short and well-conditioned.
    """

    diag = jnp.asarray(diag)
    off_diag = jnp.asarray(off_diag)
    if diag.ndim != 1 or off_diag.shape != (diag.shape[0] - 1,):
        raise ValueError(
            "diag must be 1-D and off_diag must have shape (len(diag) - 1,)"
        )
    if parallel is None:
        parallel = jax.default_backend() != "cpu" and diag.dtype == jnp.float64
    zero = jnp.zeros((1,), dtype=diag.dtype)
    lower = jnp.concatenate([zero, off_diag])
    upper = jnp.concatenate([off_diag, zero])

    # Bidiagonal Cholesky T = C C': c_off[i] = off[i] / c_d[i],
    # c_d[i+1] = sqrt(diag[i+1] - c_off[i]^2), i.e. the pivot recurrence
    # delta_{i+1} = diag_{i+1} - off_i^2 / delta_i with delta = c_d^2.
    if off_diag.shape[0] == 0:
        c_d = jnp.sqrt(diag)
        c_off = off_diag
    elif parallel:
        # The pivot recurrence is a Mobius map delta -> (a delta + b)/delta,
        # so cumulative 2x2 matrix products compose it associatively; the
        # per-element normalization leaves the projective action unchanged
        # while keeping the entries O(1).
        mobius = jnp.stack(
            [
                jnp.stack([diag[1:], -(off_diag**2)], axis=-1),
                jnp.stack([jnp.ones_like(off_diag), jnp.zeros_like(off_diag)], axis=-1),
            ],
            axis=-2,
        )

        def combine(earlier, later):
            product = later @ earlier
            scale = jnp.max(jnp.abs(product), axis=(-2, -1), keepdims=True)
            return product / scale

        cumulative = jax.lax.associative_scan(combine, mobius)
        delta_rest = (cumulative[:, 0, 0] * diag[0] + cumulative[:, 0, 1]) / (
            cumulative[:, 1, 0] * diag[0] + cumulative[:, 1, 1]
        )
        delta = jnp.concatenate([diag[:1], delta_rest])
        c_d = jnp.sqrt(delta)
        c_off = off_diag / c_d[:-1]
    else:

        def chol_step(c_d_prev, inputs):
            off_i, diag_next = inputs
            c_off_i = off_i / c_d_prev
            c_d_i = jnp.sqrt(diag_next - c_off_i**2)
            return c_d_i, (c_off_i, c_d_i)

        c_d_0 = jnp.sqrt(diag[0])
        _, (c_off, c_d_rest) = jax.lax.scan(chol_step, c_d_0, (off_diag, diag[1:]))
        c_d = jnp.concatenate([c_d_0[None], c_d_rest])

    def expand(v, x):
        return v.reshape(v.shape + (1,) * (x.ndim - 1))

    def shift_up(x):
        return jnp.concatenate([x[1:], jnp.zeros_like(x[:1])])

    def shift_down(x):
        return jnp.concatenate([jnp.zeros_like(x[:1]), x[:-1]])

    def solve(x):
        # M^{-1} x = T x: a tridiagonal matvec along the leading axis.
        return (
            expand(diag, x) * x
            + expand(upper, x) * shift_up(x)
            + expand(lower, x) * shift_down(x)
        )

    def norm(x):
        # x' M x = x' T^{-1} x through one O(n) tridiagonal solve.
        y = jax.lax.linalg.tridiagonal_solve(lower, diag, upper, x[:, None])
        return jnp.sqrt(x @ y[:, 0])

    def inv_sqrt(x):
        # S = C (lower bidiagonal) with S S' = M^{-1}.
        return expand(c_d, x) * x + expand(
            jnp.concatenate([zero, c_off]), x
        ) * shift_down(x)

    def inv_sqrt_transpose(x):
        return expand(c_d, x) * x + expand(
            jnp.concatenate([c_off, zero]), x
        ) * shift_up(x)

    return Metric(
        solve=solve,
        norm=norm,
        inv_sqrt=inv_sqrt,
        inv_sqrt_transpose=inv_sqrt_transpose,
    )

nlls_gram.metric_from_state_space(points, h, Pinf, transition, nugget=0.0, parallel=None)

Build an O(n) Metric for a stationary state-space kernel Gram.

A stationary Gaussian process has an exact O(n) Gram factorization precisely when it admits a finite-dimensional state-space (linear SDE) representation: an m-dimensional latent Gauss-Markov state observed through a row vector. Given the observation row h (m,), stationary state covariance Pinf (m, m), and transition(dt) mapping gaps (n,) to stacked transition matrices (n, m, m), the metric is M = K + nugget * I with

K[i, j] = h @ Pinf @ A[i] @ A[i-1] @ ... @ A[j+1] @ h for i > j, K[i, i] = h @ Pinf @ h, symmetric, where A[k] = transition(points[k] - points[k-1]). transition must return the TRANSPOSE of the textbook matrix exponential expm(F dt) of the SDE drift (the tinygp orientation), and transition(0) must be the identity. The main application is the half-integer Matern family — matern_state_space(sigma, ell, nu) supplies the exact (h, Pinf, transition) mapping — but sums of exponentials and other CARMA/celerite-style kernels reduce to the same form.

points must be 1-D and sorted STRICTLY increasing — not validated, since it may be traced; unsorted or repeated points silently produce a wrong or NaN metric. h, Pinf, and points may be traced (e.g. hyperparameter sweeps under jax.grad), in which case the one-time sequential Cholesky setup lands on the hot path; see the tuning guide. The ABSOLUTE nugget folds into the diagonal before factorization and is part of the metric; parallel is forwarded to metric_from_quasiseparable.

Source code in src/nlls_gram/metrics.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def metric_from_state_space(points, h, Pinf, transition, nugget=0.0, parallel=None):
    """Build an O(n) ``Metric`` for a stationary state-space kernel Gram.

    A stationary Gaussian process has an exact O(n) Gram factorization
    precisely when it admits a finite-dimensional state-space (linear SDE)
    representation: an m-dimensional latent Gauss-Markov state observed
    through a row vector. Given the observation row ``h`` (m,), stationary
    state covariance ``Pinf`` (m, m), and ``transition(dt)`` mapping gaps
    (n,) to stacked transition matrices (n, m, m), the metric is
    ``M = K + nugget * I`` with

    ``K[i, j] = h @ Pinf @ A[i] @ A[i-1] @ ... @ A[j+1] @ h`` for ``i > j``,
    ``K[i, i] = h @ Pinf @ h``, symmetric, where
    ``A[k] = transition(points[k] - points[k-1])``. ``transition`` must
    return the TRANSPOSE of the textbook matrix exponential ``expm(F dt)``
    of the SDE drift (the tinygp orientation), and ``transition(0)`` must be
    the identity. The main application is the half-integer Matern family —
    ``matern_state_space(sigma, ell, nu)`` supplies the exact
    ``(h, Pinf, transition)`` mapping — but sums of exponentials and other
    CARMA/celerite-style kernels reduce to the same form.

    ``points`` must be 1-D and sorted STRICTLY increasing — not validated,
    since it may be traced; unsorted or repeated points silently produce a
    wrong or NaN metric. ``h``, ``Pinf``, and ``points`` may be traced (e.g.
    hyperparameter sweeps under ``jax.grad``), in which case the one-time
    sequential Cholesky setup lands on the hot path; see the tuning guide.
    The ABSOLUTE ``nugget`` folds into the diagonal before factorization and
    is part of the metric; ``parallel`` is forwarded to
    ``metric_from_quasiseparable``.
    """

    d, p, q, A = quasiseparable.state_space_generators(points, h, Pinf, transition)
    return metric_from_quasiseparable(d, p, q, A, nugget=nugget, parallel=parallel)

nlls_gram.matern_state_space(sigma, ell, nu)

Return the (h, Pinf, transition) state-space model of a Matern kernel.

The half-integer Matern kernels k(tau) = sigma^2 matern_nu(tau / ell) for nu in {0.5, 1.5, 2.5} (static) are exactly the covariances of continuous-time autoregressive (CAR(m)) state-space models with state dimension m = 1, 2, 3. This returns their observation row h (m,) (sigma lives in h, not Pinf), stationary state covariance Pinf (m, m), and transition(dt) mapping gaps (n,) to stacked transition matrices (n, m, m) in the tinygp orientation (transpose of the textbook expm(F dt)), for metric_from_state_space.

With f = sqrt(2 nu) / ell: nu=0.5 has h = [sigma], Pinf = 1, transition(dt) = exp(-f dt); nu=1.5 has h = [sigma, 0], Pinf = diag(1, f^2); nu=2.5 has h = [sigma, 0, 0] and the CAR(3) Pinf/transition transcribed from tinygp v0.3.1.

Nugget-free Matern-3/2 and 5/2 Grams on fine grids are extremely ill-conditioned (condition number ~1e21 at n = 5000 — a property of the matrix, not the solver); pass an ABSOLUTE nugget (e.g. 1e-8 * sigma**2 in float64) to metric_from_state_space, which folds it into the metric exactly. For nu=0.5 the Gram inverse is exactly tridiagonal and metric_from_tridiagonal_precision is the specialized alternative whose applies are elementwise shifts.

Source code in src/nlls_gram/quasiseparable.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def matern_state_space(sigma, ell, nu):
    """Return the ``(h, Pinf, transition)`` state-space model of a Matern kernel.

    The half-integer Matern kernels
    ``k(tau) = sigma^2 matern_nu(tau / ell)`` for ``nu`` in
    ``{0.5, 1.5, 2.5}`` (static) are exactly the covariances of
    continuous-time autoregressive (CAR(m)) state-space models with state
    dimension m = 1, 2, 3. This returns their observation row ``h`` (m,)
    (``sigma`` lives in ``h``, not ``Pinf``), stationary state covariance
    ``Pinf`` (m, m), and ``transition(dt)`` mapping gaps (n,) to stacked
    transition matrices (n, m, m) in the tinygp orientation (transpose of
    the textbook ``expm(F dt)``), for ``metric_from_state_space``.

    With ``f = sqrt(2 nu) / ell``: nu=0.5 has ``h = [sigma]``, ``Pinf = 1``,
    ``transition(dt) = exp(-f dt)``; nu=1.5 has ``h = [sigma, 0]``,
    ``Pinf = diag(1, f^2)``; nu=2.5 has ``h = [sigma, 0, 0]`` and the CAR(3)
    ``Pinf``/transition transcribed from tinygp v0.3.1.

    Nugget-free Matern-3/2 and 5/2 Grams on fine grids are extremely
    ill-conditioned (condition number ~1e21 at n = 5000 — a property of the
    matrix, not the solver); pass an ABSOLUTE nugget (e.g.
    ``1e-8 * sigma**2`` in float64) to ``metric_from_state_space``, which
    folds it into the metric exactly. For nu=0.5 the Gram inverse is exactly
    tridiagonal and ``metric_from_tridiagonal_precision`` is the specialized
    alternative whose applies are elementwise shifts.
    """

    if nu not in (0.5, 1.5, 2.5):
        raise ValueError("nu must be one of 0.5, 1.5, or 2.5")
    sigma = jnp.asarray(sigma)
    ell = jnp.asarray(ell)
    one = jnp.ones_like(ell)
    zero = jnp.zeros_like(ell)

    if nu == 0.5:
        h = sigma[None]
        Pinf = one[None, None]

        def transition(dt):
            return jnp.exp(-dt / ell)[:, None, None]

    elif nu == 1.5:
        f = jnp.sqrt(3.0) / ell
        h = jnp.stack([sigma, jnp.zeros_like(sigma)])
        Pinf = jnp.stack([jnp.stack([one, zero]), jnp.stack([zero, f**2])])

        def transition(dt):
            fd = f * dt
            return jnp.exp(-fd)[:, None, None] * jnp.stack(
                [
                    jnp.stack([1.0 + fd, -(f**2) * dt], axis=-1),
                    jnp.stack([dt, 1.0 - fd], axis=-1),
                ],
                axis=-2,
            )

    else:
        f = jnp.sqrt(5.0) / ell
        f2 = f**2
        h = jnp.stack([sigma, jnp.zeros_like(sigma), jnp.zeros_like(sigma)])
        Pinf = jnp.stack(
            [
                jnp.stack([one, zero, -f2 / 3.0]),
                jnp.stack([zero, f2 / 3.0, zero]),
                jnp.stack([-f2 / 3.0, zero, f2**2]),
            ]
        )

        def transition(dt):
            fd = f * dt
            d2 = dt**2
            return jnp.exp(-fd)[:, None, None] * jnp.stack(
                [
                    jnp.stack(
                        [
                            0.5 * f2 * d2 + fd + 1.0,
                            -0.5 * f * f2 * d2,
                            0.5 * f2 * f * dt * (fd - 2.0),
                        ],
                        axis=-1,
                    ),
                    jnp.stack(
                        [dt * (fd + 1.0), -f2 * d2 + fd + 1.0, f2 * dt * (fd - 3.0)],
                        axis=-1,
                    ),
                    jnp.stack(
                        [
                            0.5 * d2,
                            0.5 * dt * (2.0 - fd),
                            0.5 * f2 * d2 - 2.0 * fd + 1.0,
                        ],
                        axis=-1,
                    ),
                ],
                axis=-2,
            )

    return h, Pinf, transition

nlls_gram.metric_from_quasiseparable(d, p, q, A, nugget=0.0, parallel=None)

Build an O(n) Metric from a symmetric quasiseparable SPD matrix.

The metric is M = K + nugget * I with K given by rank-m quasiseparable generators d (n,), p (n, m), q (n, m), and A (n, m, m):

K[i, j] = p[i] @ A[i-1] @ ... @ A[j+1] @ q[j] for i > j, d on the diagonal, and K[j, i] = K[i, j]. A[k] is the transition INTO index k (A[0] never enters the products; the state-space builders set it to the identity). nugget is an ABSOLUTE variance added to d before factorization and is part of the metric: all four callbacks act on M, consistently. Every callback costs O(n m^2) through one quasiseparable Cholesky computed at construction; solve, inv_sqrt, and inv_sqrt_transpose accept vector and matrix inputs (norm is vector-only, as everywhere in Metric).

M must be positive definite — not validated, since inputs may be traced; a non-PD or near-singular matrix silently produces NaN through the Cholesky square roots (same convention as metric_from_tridiagonal_precision).

parallel picks how the applies run, resolved ONCE at construction (a metric built on one backend keeps that scan choice for its lifetime): None (the default) uses associative O(log n)-depth scans off-CPU in float64 — where a sequential scan pays a kernel launch per step — and sequential scans otherwise. The parallel substitutions propagate rank-1-corrected transition matrices whose products are not contractive in general, so the float32 default stays sequential; pass parallel=True to override. The one-time Cholesky setup is always a sequential scan in this release.

Source code in src/nlls_gram/metrics.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
def metric_from_quasiseparable(d, p, q, A, nugget=0.0, parallel=None):
    """Build an O(n) ``Metric`` from a symmetric quasiseparable SPD matrix.

    The metric is ``M = K + nugget * I`` with ``K`` given by rank-``m``
    quasiseparable generators ``d`` (n,), ``p`` (n, m), ``q`` (n, m), and
    ``A`` (n, m, m):

    ``K[i, j] = p[i] @ A[i-1] @ ... @ A[j+1] @ q[j]`` for ``i > j``, ``d``
    on the diagonal, and ``K[j, i] = K[i, j]``. ``A[k]`` is the transition
    INTO index ``k`` (``A[0]`` never enters the products; the state-space
    builders set it to the identity). ``nugget`` is an ABSOLUTE variance
    added to ``d`` before factorization and is part of the metric: all four
    callbacks act on ``M``, consistently. Every callback costs O(n m^2)
    through one quasiseparable Cholesky computed at construction;
    ``solve``, ``inv_sqrt``, and ``inv_sqrt_transpose`` accept vector and
    matrix inputs (``norm`` is vector-only, as everywhere in ``Metric``).

    ``M`` must be positive definite — not validated, since inputs may be
    traced; a non-PD or near-singular matrix silently produces NaN through
    the Cholesky square roots (same convention as
    ``metric_from_tridiagonal_precision``).

    ``parallel`` picks how the applies run, resolved ONCE at construction
    (a metric built on one backend keeps that scan choice for its
    lifetime): ``None`` (the default) uses associative O(log n)-depth scans
    off-CPU in float64 — where a sequential scan pays a kernel launch per
    step — and sequential scans otherwise. The
    parallel substitutions propagate rank-1-corrected transition matrices
    whose products are not contractive in general, so the float32 default
    stays sequential; pass ``parallel=True`` to override. The one-time
    Cholesky setup is always a sequential scan in this release.
    """

    d = jnp.asarray(d)
    p = jnp.asarray(p)
    q = jnp.asarray(q)
    A = jnp.asarray(A)
    if d.ndim != 1:
        raise ValueError("d must be 1-D")
    n = d.shape[0]
    if p.ndim != 2 or p.shape[0] != n:
        raise ValueError("p must have shape (len(d), m)")
    m = p.shape[1]
    if q.shape != (n, m) or A.shape != (n, m, m):
        raise ValueError("q must have shape (n, m) and A shape (n, m, m)")
    d = d + nugget
    if parallel is None:
        parallel = jax.default_backend() != "cpu" and d.dtype == jnp.float64
    c, w = quasiseparable.cholesky(d, p, q, A)

    def solve(x):
        # M^{-1} x = L^{-T} L^{-1} x: forward then backward substitution.
        y = quasiseparable.forward_substitution(c, p, w, A, x.reshape(n, -1), parallel)
        return quasiseparable.backward_substitution(c, p, w, A, y, parallel).reshape(
            x.shape
        )

    def norm(x):
        # x' M x through one quasiseparable matvec, no solve.
        y = quasiseparable.matvec(d, p, q, A, x.reshape(n, -1), parallel)
        return jnp.sqrt(x @ y.reshape(x.shape))

    def inv_sqrt(x):
        # S = L^{-T} with S S' = M^{-1}.
        return quasiseparable.backward_substitution(
            c, p, w, A, x.reshape(n, -1), parallel
        ).reshape(x.shape)

    def inv_sqrt_transpose(x):
        return quasiseparable.forward_substitution(
            c, p, w, A, x.reshape(n, -1), parallel
        ).reshape(x.shape)

    return Metric(
        solve=solve,
        norm=norm,
        inv_sqrt=inv_sqrt,
        inv_sqrt_transpose=inv_sqrt_transpose,
    )

nlls_gram.metric_from_shifted_matvec(matvec, shift, *, tol=None, atol=0.0, maxiter=None, preconditioner=None)

Build a matrix-free Metric for M = A + shift * I from a matvec.

matvec(x) applies a symmetric positive-semidefinite A and must accept both (n,) vectors and (n, k) matrices on the leading axis (the same shape contract Metric.solve itself carries). shift must be POSITIVE -- it is the spectral floor that makes an iterative metric solve viable at all: cond(A + shift I) <= (lambda_max + shift) / shift regardless of how singular A is, and solve costs ~sqrt(lambda_max / shift) * log(1/tol) matvecs. A concrete shift <= 0 raises; a traced shift is not validated and must be positive. Only solve and norm are provided (there is no matrix-free square root), so the solver accepts this metric for the cholesky and cg paths and rejects qr at construction.

solve runs jax.scipy.sparse.linalg.cg on the shifted matvec, so it meets the library's exactness contract only in the tight-tol limit. Unlike dual_preconditioner error -- which never moves the converged root -- this solve's residual error perturbs the selected step, the converged solution, and the implicit derivatives at order tol; the implicit derivative in particular is computed outside the accept/reject loop with no safeguard. tol is therefore an accuracy knob for the ANSWER, not a solver schedule. The default (None -> the square root of the input dtype's machine epsilon, ~1.5e-8 in float64) keeps that perturbation at the level of typical outer tolerances; CG's attainable residual stagnates near machine_eps * (lambda_max + shift) / shift, which bounds how far tol can usefully be tightened (~1e-10 is reachable in float64 only when that condition number is ~1e4 or better). maxiter=None runs to tolerance -- a truncated CG is not a linear function of its right-hand side, which breaks the linear-operator assumptions of the cg linear_solver, so do not cap it as a cost control. preconditioner is forwarded to the inner CG (its M argument). A matrix right-hand side is solved as one stacked CG call (same spectrum, batched matvecs).

Because the inner CG is built on lax.custom_linear_solve with a symmetric operator, solve composes with jax.jvp/jax.vjp/ jax.grad — including both implicit-AD rules (dense and implicit_solver="cg") and differentiating update through the cg path — and with values the matvec closes over; matvec itself only needs to be linear. (Closure-value differentiation holds for direct Metric use and through update; the LM solve() entry point supports differentiation with respect to p only — its implicit rule is a jax.custom_jvp, which cannot close over active tracers.) Raw jax.linear_transpose of solve is not supported by JAX's CG; the solver never applies it to metric.solve (the cg implicit rule declares the application self-adjoint instead of transposing it).

Source code in src/nlls_gram/metrics.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
def metric_from_shifted_matvec(
    matvec, shift, *, tol=None, atol=0.0, maxiter=None, preconditioner=None
):
    """Build a matrix-free ``Metric`` for ``M = A + shift * I`` from a matvec.

    ``matvec(x)`` applies a symmetric positive-semidefinite ``A`` and must
    accept both ``(n,)`` vectors and ``(n, k)`` matrices on the leading axis
    (the same shape contract ``Metric.solve`` itself carries). ``shift``
    must be POSITIVE -- it is the spectral floor that makes an iterative
    metric solve viable at all: ``cond(A + shift I) <= (lambda_max + shift)
    / shift`` regardless of how singular ``A`` is, and ``solve`` costs
    ``~sqrt(lambda_max / shift) * log(1/tol)`` matvecs. A concrete
    ``shift <= 0`` raises; a traced ``shift`` is not validated and must be
    positive. Only ``solve`` and ``norm`` are provided (there is no
    matrix-free square root), so the solver accepts this metric for the
    ``cholesky`` and ``cg`` paths and rejects ``qr`` at construction.

    ``solve`` runs ``jax.scipy.sparse.linalg.cg`` on the shifted matvec, so
    it meets the library's exactness contract only in the tight-``tol``
    limit. Unlike ``dual_preconditioner`` error -- which never moves the
    converged root -- this solve's residual error perturbs the selected
    step, the converged solution, and the implicit derivatives at order
    ``tol``; the implicit derivative in particular is computed outside the
    accept/reject loop with no safeguard. ``tol`` is therefore an accuracy
    knob for the ANSWER, not a solver schedule. The default (``None`` ->
    the square root of the input dtype's machine epsilon, ~1.5e-8 in
    float64) keeps that perturbation at the level of typical outer
    tolerances; CG's attainable residual stagnates near
    ``machine_eps * (lambda_max + shift) / shift``, which bounds how far
    ``tol`` can usefully be tightened (~1e-10 is reachable in float64 only
    when that condition number is ~1e4 or better). ``maxiter=None`` runs to
    tolerance -- a truncated CG is not a linear function of its right-hand
    side, which breaks the linear-operator assumptions of the ``cg``
    linear_solver, so do not cap it as a cost control. ``preconditioner``
    is forwarded to the inner CG (its ``M`` argument). A matrix right-hand
    side is solved as one stacked CG call (same spectrum, batched matvecs).

    Because the inner CG is built on ``lax.custom_linear_solve`` with a
    symmetric operator, ``solve`` composes with ``jax.jvp``/``jax.vjp``/
    ``jax.grad`` — including both implicit-AD rules (dense and
    ``implicit_solver="cg"``) and differentiating ``update`` through the
    ``cg`` path — and with values the matvec closes over; ``matvec`` itself
    only needs to be linear. (Closure-value differentiation holds for
    direct ``Metric`` use and through ``update``; the LM ``solve()``
    entry point supports differentiation with respect to ``p`` only — its
    implicit rule is a ``jax.custom_jvp``, which cannot close over active
    tracers.) Raw ``jax.linear_transpose`` of ``solve`` is not supported by
    JAX's CG; the solver never applies it to ``metric.solve`` (the cg
    implicit rule declares the application self-adjoint instead of
    transposing it).
    """

    if not isinstance(shift, jax.core.Tracer) and float(shift) <= 0.0:
        raise ValueError("shift must be positive (it is the spectral floor)")
    if tol is not None and tol < 0:
        raise ValueError("tol must be nonnegative")
    if atol < 0:
        raise ValueError("atol must be nonnegative")
    if maxiter is not None and maxiter <= 0:
        raise ValueError("maxiter must be positive or None")

    def shifted_matvec(x):
        return matvec(x) + shift * x

    def solve(x):
        if tol is None:
            solve_tol = float(jnp.finfo(jnp.result_type(x)).eps) ** 0.5
        else:
            solve_tol = tol
        solution, _ = jsp_sparse_linalg.cg(
            shifted_matvec,
            x,
            tol=solve_tol,
            atol=atol,
            maxiter=maxiter,
            M=preconditioner,
        )
        return solution

    def norm(x):
        return jnp.sqrt(x @ shifted_matvec(x))

    return Metric(solve=solve, norm=norm)

nlls_gram.metric_from_diagonal(weights)

Build a Metric for the diagonal metric M = diag(weights).

weights (length n) must be positive -- not validated, since inputs may be traced. Every callback is elementwise.

Source code in src/nlls_gram/metrics.py 
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
def metric_from_diagonal(weights):
    """Build a ``Metric`` for the diagonal metric ``M = diag(weights)``.

    ``weights`` (length n) must be positive -- not validated, since inputs
    may be traced. Every callback is elementwise.
    """

    weights = jnp.asarray(weights)
    if weights.ndim != 1:
        raise ValueError("weights must be 1-D")
    sqrt_weights = jnp.sqrt(weights)

    def expand(v, x):
        return v.reshape(v.shape + (1,) * (x.ndim - 1))

    def solve(x):
        return x / expand(weights, x)

    def norm(x):
        return jnp.sqrt(x @ (weights * x))

    def inv_sqrt(x):
        return x / expand(sqrt_weights, x)

    return Metric(
        solve=solve,
        norm=norm,
        inv_sqrt=inv_sqrt,
        inv_sqrt_transpose=inv_sqrt,
    )

nlls_gram.blockdiag_metric(blocks)

Compose per-block Metrics into one block-diagonal Metric.

blocks is a sequence of (metric, size) pairs in the order the flattened parameter vector is laid out. solve, inv_sqrt, and inv_sqrt_transpose slice on the leading axis (vector and matrix inputs both work); norm combines the block norms in quadrature.

A fully-default Metric() block means the identity metric on that block, and its callbacks are filled in accordingly. A block that defines some callbacks but leaves others None is treated as missing those callbacks: the composite field is None, so the solver's construction-time validation applies exactly as it would to that block alone. Filling identity there instead would silently break the S S' = M^{-1} consistency between callbacks.

Source code in src/nlls_gram/metrics.py 
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
def blockdiag_metric(blocks):
    """Compose per-block ``Metric``s into one block-diagonal ``Metric``.

    ``blocks`` is a sequence of ``(metric, size)`` pairs in the order the
    flattened parameter vector is laid out. ``solve``, ``inv_sqrt``, and
    ``inv_sqrt_transpose`` slice on the leading axis (vector and matrix
    inputs both work); ``norm`` combines the block norms in quadrature.

    A fully-default ``Metric()`` block means the identity metric on that
    block, and its callbacks are filled in accordingly. A block that defines
    some callbacks but leaves others ``None`` is treated as missing those
    callbacks: the composite field is ``None``, so the solver's
    construction-time validation applies exactly as it would to that block
    alone. Filling identity there instead would silently break the
    ``S S' = M^{-1}`` consistency between callbacks.
    """

    if not blocks:
        raise ValueError("blocks must contain at least one (metric, size) pair")
    metrics = [metric for metric, _ in blocks]
    offsets = [0]
    for _, size in blocks:
        offsets.append(offsets[-1] + size)

    identity_block = [
        metric.solve is None
        and metric.norm is None
        and metric.inv_sqrt is None
        and metric.inv_sqrt_transpose is None
        for metric in metrics
    ]

    def split(x):
        return [
            x[start:stop] for start, stop in zip(offsets[:-1], offsets[1:], strict=True)
        ]

    def identity(x):
        return x

    def blockwise(callbacks):
        if any(
            callback is None and not is_identity
            for callback, is_identity in zip(callbacks, identity_block, strict=True)
        ):
            return None
        filled = [identity if callback is None else callback for callback in callbacks]

        def apply(x):
            return jnp.concatenate(
                [
                    callback(part)
                    for callback, part in zip(filled, split(x), strict=True)
                ]
            )

        return apply

    def make_norm(callbacks):
        if any(
            callback is None and not is_identity
            for callback, is_identity in zip(callbacks, identity_block, strict=True)
        ):
            return None
        filled = [
            jnp.linalg.norm if callback is None else callback for callback in callbacks
        ]

        def norm(x):
            parts = split(x)
            return jnp.sqrt(
                sum(
                    callback(part) ** 2
                    for callback, part in zip(filled, parts, strict=True)
                )
            )

        return norm

    return Metric(
        solve=blockwise([metric.solve for metric in metrics]),
        norm=make_norm([metric.norm for metric in metrics]),
        inv_sqrt=blockwise([metric.inv_sqrt for metric in metrics]),
        inv_sqrt_transpose=blockwise([metric.inv_sqrt_transpose for metric in metrics]),
    )

nlls_gram.sherman_morrison_preconditioner(solve, u, weight)

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. 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 damping argument is accepted per the dual_preconditioner contract and ignored -- spectral closeness to the damped operator is all a preconditioner needs -- which also makes the helper directly valid as implicit_preconditioner (the solver calls two-argument callables with zero damping there).

Source code in src/nlls_gram/preconditioners.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def sherman_morrison_preconditioner(solve, u, weight):
    """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. 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 ``damping`` argument is accepted per the
    ``dual_preconditioner`` contract and ignored -- spectral closeness to the
    damped operator is all a preconditioner needs -- which also makes the
    helper directly valid as ``implicit_preconditioner`` (the solver calls
    two-argument callables with zero damping there).
    """

    solve_u = solve(u)
    denominator = 1.0 / weight + u @ solve_u

    def dual_preconditioner(v, damping):
        y = solve(v)
        return y - solve_u * ((u @ y) / denominator)

    return dual_preconditioner

nlls_gram.woodbury_preconditioner(solve, U, weights)

Preconditioner for B = A + U diag(weights) U' from a solve with A.

The rank-k generalization of sherman_morrison_preconditioner: 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 by the Woodbury identity; A^{-1} U (one matrix solve) and the Cholesky factor of the k x k capacitance are precomputed. This is the natural shape when a metric weight eps on a k-vector of scalar parameters injects the exactly known rank-k spike (c^2/eps) U U' into J M^{-1} J' (U the corresponding Jacobian columns up to sign and scale). With k = 1 it reduces to sherman_morrison_preconditioner. weights must be positive -- not validated, since inputs may be traced. The damping argument is accepted per the dual_preconditioner contract and ignored, so the helper is directly valid as implicit_preconditioner too.

Source code in src/nlls_gram/preconditioners.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def woodbury_preconditioner(solve, U, weights):
    """Preconditioner for ``B = A + U diag(weights) U'`` from a solve with ``A``.

    The rank-k generalization of ``sherman_morrison_preconditioner``:
    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`` by the Woodbury
    identity; ``A^{-1} U`` (one matrix solve) and the Cholesky factor of the
    k x k capacitance are precomputed. This is the natural shape when a
    metric weight ``eps`` on a k-vector of scalar parameters injects the
    exactly known rank-k spike ``(c^2/eps) U U'`` into ``J M^{-1} J'``
    (``U`` the corresponding Jacobian columns up to sign and scale). With
    ``k = 1`` it reduces to ``sherman_morrison_preconditioner``. ``weights``
    must be positive -- not validated, since inputs may be traced. The
    ``damping`` argument is accepted per the ``dual_preconditioner``
    contract and ignored, so the helper is directly valid as
    ``implicit_preconditioner`` too.
    """

    U = jnp.asarray(U)
    weights = jnp.asarray(weights)
    if U.ndim != 2 or weights.shape != (U.shape[1],):
        raise ValueError("U must have shape (n, k) and weights shape (k,)")
    solve_U = solve(U)
    capacitance = jnp.diag(1.0 / weights) + U.T @ solve_U
    factor = jsp_linalg.cho_factor(capacitance)

    def dual_preconditioner(v, damping):
        y = solve(v)
        return y - solve_U @ jsp_linalg.cho_solve(factor, U.T @ y)

    return dual_preconditioner

nlls_gram.identity_preconditioner()

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

linear_solver="cg" requires dual_preconditioner, and a cg-resolved implicit solve requires implicit_preconditioner -- nobody should run Krylov methods without thinking about preconditioning, so opting out is an explicit, greppable decision rather than a silent default. The returned callable accepts both hook signatures: dual_preconditioner(v, damping) and implicit_preconditioner(v).

Source code in src/nlls_gram/preconditioners.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def identity_preconditioner():
    """The identity map as an explicit "no preconditioner" choice.

    ``linear_solver="cg"`` requires ``dual_preconditioner``, and a cg-resolved
    implicit solve requires ``implicit_preconditioner`` -- nobody should run
    Krylov methods without thinking about preconditioning, so opting out is an
    explicit, greppable decision rather than a silent default. The returned
    callable accepts both hook signatures: ``dual_preconditioner(v, damping)``
    and ``implicit_preconditioner(v)``.
    """

    def preconditioner(v, damping=None):
        return v

    return preconditioner

nlls_gram.nystrom_preconditioner(matvec, n, rank, key, *, dtype=None)

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 Frangella, Tropp, and Udell (arXiv:2110.02820, Algorithm 2.1); the stabilization shift nu ~ eps * ||Y||_F is removed from the recovered eigenvalues. The returned callback applies the FTU preconditioner (their eq. 5.3, up to the positive scalar rho + damping, which CG ignores)::

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

where rho is the smallest retained Nystrom eigenvalue: eigendirections 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. This is the one shipped helper that uses the live damping argument (Sherman-Morrison/Woodbury ignore it): one construction serves every LM damping value, and passed as implicit_preconditioner it is called with zero damping and applies the undamped inverse (valid only when the retained spectrum is strictly positive).

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 and accept (n, k) matrices (the same shape contract as Metric.solve); an indefinite operator silently produces NaN through the Cholesky square root. The build costs rank operator applications plus an O(n rank^2) QR/SVD, done once at construction -- like every preconditioner it is frozen there, so for a nonlinear problem it approximates the dual at the linearization point it was built from (staleness is safe: preconditioner error never moves the converged root). 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 (respects x64) -- pass the operator dtype explicitly for a float32 problem under enabled x64. All operations are traceable; n and rank are static Python ints.

Source code in src/nlls_gram/preconditioners.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def nystrom_preconditioner(matvec, n, rank, key, *, dtype=None):
    """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 Frangella, Tropp, and Udell (arXiv:2110.02820, Algorithm 2.1); the
    stabilization shift ``nu ~ eps * ||Y||_F`` is removed from the recovered
    eigenvalues. The returned callback applies the FTU preconditioner
    (their eq. 5.3, up to the positive scalar ``rho + damping``, which CG
    ignores)::

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

    where ``rho`` is the smallest retained Nystrom eigenvalue: eigendirections
    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. This is the one shipped helper that uses the
    live ``damping`` argument (Sherman-Morrison/Woodbury ignore it): one
    construction serves every LM damping value, and passed as
    ``implicit_preconditioner`` it is called with zero damping and applies
    the undamped inverse (valid only when the retained spectrum is strictly
    positive).

    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 and accept
    ``(n, k)`` matrices (the same shape contract as ``Metric.solve``); an
    indefinite operator silently produces NaN through the Cholesky square
    root. The build costs ``rank`` operator applications plus an
    ``O(n rank^2)`` QR/SVD, done once at construction -- like every
    preconditioner it is frozen there, so for a nonlinear problem it
    approximates the dual at the linearization point it was built from
    (staleness is safe: preconditioner error never moves the converged
    root). 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 (respects
    x64) -- pass the operator dtype explicitly for a float32 problem under
    enabled x64. All operations are traceable; ``n`` and ``rank`` are static
    Python ints.
    """

    if not isinstance(n, int) or isinstance(n, bool) or n <= 0:
        raise ValueError("n must be a positive int")
    if not isinstance(rank, int) or isinstance(rank, bool) or not 0 < rank <= n:
        raise ValueError("rank must be a positive int <= n")
    if dtype is None:
        dtype = jnp.result_type(float)
    Omega = jnp.linalg.qr(jax.random.normal(key, (n, rank), dtype=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 = 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)
    lam = jnp.maximum(sigma**2 - nu, 0.0)
    rho = lam[-1]

    def preconditioner(v, damping=0.0):
        # U (U'v)/(lam+damping) + (v - U U'v)/(rho+damping), regrouped so the
        # apply is two (n, rank) matvecs instead of three.
        Utv = U.T @ v
        return U @ (Utv / (lam + damping) - Utv / (rho + damping)) + v / (rho + damping)

    return preconditioner

nlls_gram.pad_dual_preconditioner(base_preconditioner, n_real)

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

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

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

and the matching preconditioner applies base_preconditioner on the first n_real coordinates and the exact 1 / damping inverse on the padded block -- the 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). Wrapping is needed for shape-fixed bases (dense solves, nystrom_preconditioner, Sherman-Morrison/Woodbury built at the unpadded size); a shape-generic base like identity_preconditioner() stays valid unwrapped, it just forgoes the exact padded-block inverse. Like nystrom_preconditioner this uses the live damping argument, and because the padded block divides by it, the returned callback serves only the damped forward solve -- never the implicit_preconditioner hook. Relatedly, padded rows make the undamped implicit dual J P J' singular, so the library's implicit rules return a non-finite derivative on padded problems; the minimum-metric-norm derivative still exists mathematically and equals the unpadded one, so differentiate the unpadded formulation.

Source code in src/nlls_gram/preconditioners.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def pad_dual_preconditioner(base_preconditioner, n_real):
    """Extend a dual preconditioner to a residual padded with exact zeros.

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

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

    and the matching preconditioner applies ``base_preconditioner`` on the
    first ``n_real`` coordinates and the exact ``1 / damping`` inverse on the
    padded block -- the 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). Wrapping is needed
    for shape-fixed bases (dense solves, ``nystrom_preconditioner``,
    Sherman-Morrison/Woodbury built at the unpadded size); a shape-generic
    base like ``identity_preconditioner()`` stays valid unwrapped, it just
    forgoes the exact padded-block inverse. Like ``nystrom_preconditioner``
    this uses the live ``damping`` argument, and because the padded block
    divides by it, the returned callback serves only the damped forward
    solve -- never the ``implicit_preconditioner`` hook. Relatedly, padded
    rows make the undamped implicit dual ``J P J'`` singular, so the
    library's implicit rules return a non-finite derivative on padded
    problems; the minimum-metric-norm derivative still exists mathematically
    and equals the unpadded one, so differentiate the unpadded formulation.
    """

    if not isinstance(n_real, int) or isinstance(n_real, bool) or n_real <= 0:
        raise ValueError("n_real must be a positive int")

    def dual_preconditioner(v, damping):
        # 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] < n_real:
            raise ValueError(
                f"padded residual vector must be 1-D with at least "
                f"n_real={n_real} entries; got shape {v.shape}"
            )
        return jnp.concatenate(
            (base_preconditioner(v[:n_real], damping), v[n_real:] / damping)
        )

    # The padded block divides by the live damping, so the zero-damping
    # implicit hook must reject this helper at construction.
    dual_preconditioner.requires_positive_damping = True
    return dual_preconditioner