Make QuasisepSolver.condition respect parallel=True and use a rank-J conditional covariance

[dfm]

Fixes #280.

Summary

parallel=True on QuasisepSolver did not reach condition. Conditioning at the training coordinates ran four sequential scans: three inside qsm_mul (one for L^{-1} @ M, two for the .gram()), and one in the Cholesky of the conditioned GaussianProcess, which never inherited the parallel flag. The issue's timings on GPU line up with exactly that count.

This PR fixes it in two steps.

1. Parallel qsm_mul, and forward the flag

The two lax.scan recurrences in qsm_mul are affine maps with a two-sided linear part, phi -> a @ phi @ b.T + B. Those compose without any Kronecker product by carrying triples (a, b, B):

def combine(left, right):
    (al, bl, Bl), (ar, br, Br) = left, right
    return ar @ al, br @ bl, ar @ Bl @ br.T + Br

so both scans now have associative_scan versions, and qsm_mul and SquareQSM.gram take a static parallel keyword. gp.condition forwards parallel to the conditioned GP when the parent solver is a QuasisepSolver and the covariance came back as a SymmQSM, and QuasisepSolver.covariance respects the flag as well.

2. A rank-J representation for the same-kernel conditional

An adversarial review of step 1 found a real regression: forwarding parallel meant the conditioned GP's covariance was factorized by the existing cholesky_parallel, which lost several digits on it and produced NaN samples and -inf log probabilities in ordinary configurations (e.g. Matern32, N=1000, diag=1e-2, default conditioning jitter).

The cause is the representation, not the condition number. The conditional at the training points was built as M - (L^{-1} M)^T (L^{-1} M), which has 4J generators that encode the difference of two nearly equal matrices. When conditioning with the GP's own kernel, the same matrix is N - N @ K^{-1} @ N with N the noise, which has rank J and needs only symm_inv (which already has a parallel version). On that representation the parallel Cholesky error drops from NaN to ~1e-8, and even the sequential Cholesky improves by four to six orders of magnitude:

kernel, N	rank	seq Cholesky err	par Cholesky err
Matern32, 500 (old form)	8	4.0e-06	nan
Matern32, 500 (new form)	2	6.7e-12	2.2e-10
Cosine, 500 (old form)	8	4.6e-06	nan
Cosine, 500 (new form)	2	1.0e-12	6.9e-09

To use this, kernel=None is now passed through to Solver.condition and Solver.condition_diag as the "same kernel" signal (the same convention #272 adopts), both built-in solvers store self.kernel, and QuasisepSolver stores its noise. The cross-kernel path keeps the product form, now with parallel scans.

Along the way this also fixes a pre-existing bug where qsm_mul labelled the product of two different SymmQSMs as symmetric and dropped the upper triangle.

Verification

• Full suite passes. New tests cover qsm_mul parity with dense and Block transition matrices, non-symmetric products of symmetric operands, Banded noise conditioning, chained conditioning, predict(return_var=True), and the conditioned GP's log_probability and sample in both modes under x64 for Matern32, Cosine, and a sum kernel. The Cosine case fails on step 1 alone.
• jax.make_jaxpr shows zero scan primitives on a parallel=True GP for log_probability, condition, the conditioned GP's log_probability, predict(return_var=True), and covariance.
• Reviewer stress grid (N=1000, parent diag from 1 to 1e-4, three conditioning jitters): zero NaN factors or samples in every cell, in both modes.
• XLA temp memory for parallel condition dropped from 3128 to 188 bytes per data point (about 1.9 GB instead of 31 GB at N=1e7 in float64).
• CPU timing from the review (N=2e5, float64, Matern32): parallel condition 0.5 s vs 90 s sequential, before the rank-J change removed the dominant term.

Notes

• Breaking for third-party Solver subclasses: condition/condition_diag may now receive kernel=None and must store self.kernel. Add O(J^2)-per-test-point predictive mean/variance from the QSM Cholesky #272 makes the same change, so the two converge semantically but will conflict textually in solver.py and gp.condition.
• Still sequential on a parallel GP, pre-existing and out of scope here: the mean with an explicit kernel= argument, and predictions at new X_test (via GeneralQSM). cholesky_parallel also remains fragile on badly represented matrices in general; probably worth its own issue.

🤖 Generated with Claude Code

https://claude.ai/code/session_01JTMk4JdA8mBFozHWnxqUqy