Add O(J^2)-per-test-point predictive variance from the QSM Cholesky

[dfm]

Stacked on #282, which should land first; the diff here is the two commits on top of it.

Summary

When a QuasisepSolver conditions at new test points with the kernel it was built with (kernel=None, the default), the predictive variance is now computed in O(J²) per test point (after two O(N J²) train-only scans) by reusing the quasiseparable Cholesky factorization, instead of a dense O(M² N) conditional covariance. The predictive mean was already scalable via kernel.matmul(X_test, X_train, alpha) and is unchanged. The conditioned GaussianProcess gets a LazyDirectSolver, so its dense covariance is only built (and factorized) when covariance, sample, or log_probability are explicitly requested.

This comes in two commits.

1. Refactor the solver conditioning interface around a single hook

Solver.condition now takes the unresolved kernel (None meaning "the kernel this solver was built with", a signal that survives jax.jit where object identity does not), the test coordinates, the test noise, and alpha = K^{-1} (y - mean), and returns a ConditionedComponents bundle with the conditional kernel, the conditional mean at the test points, and the conditioned process's solver. The producing solver builds that solver itself, so settings such as parallel carry over without special cases in GaussianProcess (this subsumes the parallel forwarding in #282's gp.py), and GaussianProcess accepts an already constructed solver instance.

GaussianProcess.predict is now a thin wrapper around condition. Under jax.jit, XLA eliminates the unused N_test x N_test covariance and its Cholesky factorization, so the separate Solver.condition_diag hook from #281 and the mean-only shortcuts are removed. (Verified from compiled HLO: on the pre-#281 commit, jitted predict(return_var=True) already compiled to zero M x M buffers and zero potrf calls for both solvers.) The generic dense implementation lives in tinygp.solvers.direct.dense_condition, which evaluates the cross covariance block once and passes the variance to the resulting DirectSolver explicitly, so the full matrix stays dead code when only the variance is read.

Also: kernels.Conditioned gains a batched __call__; means.Conditioned.include_mean is a static field (it was a traced leaf, which broke conditioning a conditioned process under jax.jit); the default jitter for the conditioned process follows the data dtype.

Breaking for third-party Solver subclasses: condition has a new signature and return type, condition_diag no longer exists, and solvers must expose kernel and X.

2. Add the O(J²)-per-test-point predictive variance

The variance uses the same per-point cross-covariance row generators as the rectangular quasiseparable product, factored out of Quasisep.to_general_qsm into Quasisep.anchor, so there is a single anchoring convention. Every propagation runs forward in time, which keeps the result stable across wide training gaps (a formulation that pulls a test point back across its gap with an inverse transition overflows for gaps of a few dozen correlation lengths). anchor evaluates masked transitions at a clamped coordinate, so gradients at far extrapolation stay finite; this also fixes a latent NaN gradient in kernel.matmul for extrapolating test points.

ops.cholesky and ops.cholesky_parallel also return the inclusive Riccati carry, and the backward congruence recursion shared with symm_inv_parallel lives in ops.congruence_scan. Noise models with their own quasiseparable states (e.g. Banded) enlarge the factorized matrix beyond the kernel's order, so those cases take the dense path.

Bugs fixed along the way

• Conditioning a conditioned process (gp.condition(y).gp.condition(y2, X_test)) failed on main with the default include_mean=True.
• main's QuasisepSolver dense fallback omitted the test noise from the conditional covariance (inconsistent with DirectSolver and with its own variance).
• NaN gradient of kernel.matmul at far extrapolation.

Verification

• Full suite passes (302 tests), including new tests for the fast path vs a dense reference across six kernel fixtures × sequential/parallel; edge geometries (coincident points, duplicated training times, N=1, N=2, extrapolation both sides, unsorted test points); wide training gaps (40 and 500 correlation lengths); banded training noise falling back to dense; NumPy inputs; the fast path surviving a jit pytree round-trip; and chained conditioning keeping its quasiseparable representation.
• Compiled-HLO probe (N=300, M=1500, float64, Matern32), counting f64[M,M] buffers / potrf / XLA temp memory for jitted predict(return_var=True): DirectSolver 0 / 0 / 11.5 MB; QuasisepSolver 0 / 0 / 0.3 MB (25.3 MB with N x M buffers before this change).
• Numerically identical to main for DirectSolver across X_test, kernel=, include_mean, diag/noise, return_var/return_cov, and mixed dtypes; QuasisepSolver matches DirectSolver to <1e-8 everywhere, including chained conditioning under jit with parallel=True.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HBWVoYn1nqDHDaF1K89Vrd