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

news/280.bugfix

@@ -0,0 +1,15 @@
+Fixed ``parallel=True`` not reaching the ``QuasisepSolver.condition`` path when
+conditioning at the training coordinates. The quasiseparable matrix products
+used to build the conditional covariance now support parallel associative
+scans, the conditioned ``GaussianProcess`` inherits the ``parallel`` flag from
+its parent, and ``QuasisepSolver.covariance`` also respects the flag. As part
+of this change, conditioning at the training coordinates with the same kernel
+that the GP was built with now uses the identity ``N - N @ K^{-1} @ N`` for the
+conditional covariance, which has the same quasiseparable rank as the prior
+kernel (instead of four times that rank) and is much better conditioned
+numerically. To support this, ``kernel=None`` is now passed through to
+``Solver.condition`` and ``Solver.condition_diag`` to signal that the solver's
+own kernel should be used, and solvers are expected to store that kernel as
+``self.kernel``; third-party ``Solver`` implementations will need to handle
+this. Also fixed a bug where the product of two different ``SymmQSM`` matrices
+was incorrectly returned as a ``SymmQSM``, dropping its upper triangle.

src/tinygp/gp.py

@@ -176,17 +176,29 @@ def condition(
         _check_test_shapes(self.X, X_test)
 
         alpha, log_prob, mean_value = self._condition(y, X_test, include_mean, kernel)
-        if kernel is None:
-            kernel = self.kernel
 
         if noise is None:
             diag = _default_diag(mean_value) if diag is None else diag
             noise = Diagonal(diag=jnp.broadcast_to(diag, mean_value.shape))
 
+        # Note: ``kernel`` is passed through unresolved here, since ``None``
+        # signals to the solver that we're conditioning using the same kernel
+        # that it was built with, and it can use a more efficient algorithm
         covariance_value = self.solver.condition(kernel, X_test, noise)
+        if kernel is None:
+            kernel = self.kernel
         if X_test is None:
             X_test = self.X
 
+        # When the conditional covariance is quasiseparable, the new GP will
+        # use a QuasisepSolver, and it should inherit the parallel flag from
+        # this GP's solver
+        solver_kwargs: dict[str, Any] = {}
+        if isinstance(self.solver, QuasisepSolver) and isinstance(
+            covariance_value, SymmQSM
+        ):
+            solver_kwargs["parallel"] = self.solver.parallel
+
         # The conditional GP will also be a GP with the mean an covariance
         # specified by a :class:`tinygp.means.Conditioned` and
         # :class:`tinygp.kernels.Conditioned` respectively.
@@ -203,6 +215,7 @@ def condition(
             ),
             mean_value=mean_value,
             covariance_value=covariance_value,
+            **solver_kwargs,
         )
 
         return ConditionResult(log_prob, gp)
@@ -273,7 +286,7 @@ def predict(
             _, _, mean_value = self._condition(y, None, include_mean, kernel)
             diag = _default_diag(mean_value)
             noise = Diagonal(diag=jnp.broadcast_to(diag, mean_value.shape))
-            var_value = self.solver.condition_diag(pred_kernel, None, noise)
+            var_value = self.solver.condition_diag(kernel, None, noise)
             return mean_value, var_value
 
         # Predicting at new test points: build `Ks` once and reuse it for

src/tinygp/solvers/direct.py

@@ -22,6 +22,7 @@ class DirectSolver(Solver):
     usual constructor.
     """
 
+    kernel: kernels.Kernel
     X: JAXArray
     variance_value: JAXArray
     covariance_value: JAXArray
@@ -45,6 +46,7 @@ def __init__(
                 matrix. This should be equal to the result of calling ``kernel``
                 and adding ``diag``, but that is not checked.
         """
+        self.kernel = kernel
         self.X = X
         self.variance_value = kernel(X) + noise.diagonal()
         if covariance is None:
@@ -73,17 +75,20 @@ def dot_triangular(self, y: JAXArray) -> JAXArray:
         return jnp.einsum("ij,j...->i...", self.scale_tril, y)
 
     def condition(
-        self, kernel: kernels.Kernel, X_test: JAXArray | None, noise: Noise
+        self, kernel: kernels.Kernel | None, X_test: JAXArray | None, noise: Noise
     ) -> Any:
         """Compute the covariance matrix for a conditional GP
 
         Args:
             kernel: The kernel for the covariance between the observed and
-                predicted data.
+                predicted data. If ``None``, the kernel used to construct this
+                solver is used.
             X_test: The coordinates of the predicted points. Defaults to the
                 input coordinates.
             noise: The noise model for the predicted process.
         """
+        if kernel is None:
+            kernel = self.kernel
         if X_test is None:
             Ks = kernel(self.X, self.X)
             Kss = Ks + noise

src/tinygp/solvers/kalman.py

@@ -79,7 +79,9 @@ def dot_triangular(self, y: JAXArray) -> JAXArray:
         del y
         raise NotImplementedError
 
-    def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any:
+    def condition(
+        self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise
+    ) -> Any:
         del kernel, X_test, noise
         raise NotImplementedError
 

src/tinygp/solvers/quasisep/core.py

@@ -421,16 +421,22 @@ def scale(self, other: JAXArray) -> SquareQSM:
             upper=self.upper.scale(other),
         )
 
-    def gram(self) -> SymmQSM:
+    def gram(self, *, parallel: bool = False) -> SymmQSM:
         """The inner product of this matrix with itself
 
         If this matrix is called ``A``, the Gram matrix is ``A.T @ A``, and
         that's what this method computes. The result is a :class:`SymmQSM`.
+
+        Args:
+            parallel: If ``True``, use parallel associative-scan algorithms
+                for the matrix product.
         """
+        from tinygp.solvers.quasisep.ops import qsm_mul
+
         # We know that this must result in symmetric matrix, but that won't be
         # enforced; we make it so! It might be possible to make this more
         # efficient, but perhaps jax is clever enough?
-        M = self.transpose() @ self
+        M = qsm_mul(self.transpose(), self, parallel=parallel)
         return SymmQSM(diag=M.diag, lower=M.lower)
 
     @jax.jit

src/tinygp/solvers/quasisep/ops.py

@@ -2,6 +2,7 @@
 
 __all__ = ["elementwise_add", "elementwise_mul", "qsm_mul"]
 
+from functools import partial
 from typing import TypeVar
 
 import jax
@@ -49,8 +50,16 @@ def elementwise_mul(a: QSM, b: QSM) -> QSM | None:
     return construct(diag, lower, upper, is_symm_a and is_symm_b)
 
 
-@jax.jit
-def qsm_mul(a: QSM, b: QSM) -> QSM | None:
+@partial(jax.jit, static_argnames=("parallel",))
+def qsm_mul(a: QSM, b: QSM, *, parallel: bool = False) -> QSM | None:
+    """The product of two quasiseparable matrices
+
+    Args:
+        a: The left matrix.
+        b: The right matrix.
+        parallel: If ``True``, use parallel associative-scan algorithms for the
+            two recurrences, instead of sequential scans.
+    """
     diag_a, lower_a, upper_a = deconstruct(a)
     diag_b, lower_b, upper_b = deconstruct(b)
 
@@ -60,28 +69,14 @@ def qsm_mul(a: QSM, b: QSM) -> QSM | None:
         return DiagQSM(d=diag_a * diag_b)
 
     if lower_a is not None and upper_b is not None:
-
-        def calc_phi(phi, data):  # type: ignore
-            a, b, q, g = data
-            return a @ phi @ b.T + jnp.outer(q, g), phi
-
-        init = jnp.zeros_like(jnp.outer(lower_a.q[0], upper_b.q[0]))
-        args = (lower_a.a, upper_b.a, lower_a.q, upper_b.q)
-        _, phi = jax.lax.scan(calc_phi, init, args)
-
+        impl = qsm_mul_phi_parallel if parallel else qsm_mul_phi
+        phi = impl(lower_a.a, upper_b.a, lower_a.q, upper_b.q)
     else:
         phi = None
 
     if upper_a is not None and lower_b is not None:
-
-        def calc_psi(psi, data):  # type: ignore
-            a, b, q, g = data
-            return a.T @ psi @ b + jnp.outer(q, g), psi
-
-        init = jnp.zeros_like(jnp.outer(upper_a.q[-1], lower_b.p[-1]))
-        args = (upper_a.a, lower_b.a, upper_a.p, lower_b.p)
-        _, psi = jax.lax.scan(calc_psi, init, args, reverse=True)
-
+        impl = qsm_mul_psi_parallel if parallel else qsm_mul_psi
+        psi = impl(upper_a.a, lower_b.a, upper_a.p, lower_b.p)
     else:
         psi = None
 
@@ -209,9 +204,11 @@ def impl(
     diag, lower, upper = impl(
         diag_a, lower_a, upper_a, diag_b, lower_b, upper_b, phi, psi
     )
-    is_symm_a = isinstance(a, (DiagQSM, SymmQSM))
-    is_symm_b = isinstance(b, (DiagQSM, SymmQSM))
-    return construct(diag, lower, upper, is_symm_a and is_symm_b)
+    # Note: the product of two symmetric matrices is not, in general,
+    # symmetric, so we always return the full (square) result here. Callers
+    # that know the result must be symmetric (e.g. ``SquareQSM.gram``) can
+    # re-wrap it.
+    return construct(diag, lower, upper, False)
 
 
 def deconstruct(
@@ -305,6 +302,52 @@ def _shift_bwd(x):
     return jnp.concatenate((x[1:], jnp.zeros_like(x[-1:])), axis=0)
 
 
+@jax.jit
+def qsm_mul_phi(a, b, q, g):
+    def impl(phi, data):
+        a, b, q, g = data
+        return a @ phi @ b.T + jnp.outer(q, g), phi
+
+    init = jnp.zeros_like(jnp.outer(q[0], g[0]))
+    _, phi = jax.lax.scan(impl, init, (a, b, q, g))
+    return phi
+
+
+@jax.jit
+def qsm_mul_phi_parallel(a, b, q, g):
+    # The recurrence is affine in phi, with a two-sided linear part:
+    # phi -> a @ phi @ b.T + B. The composition of two such maps is again of
+    # this form, so we scan over the triples (a, b, B).
+    B = jnp.einsum("nj,nk->njk", q, g)
+    _, _, phi = jax.lax.associative_scan(_two_sided_affine_combine, (a, b, B))
+    return _shift_fwd(phi)
+
+
+@jax.jit
+def qsm_mul_psi(a, b, q, g):
+    def impl(psi, data):
+        a, b, q, g = data
+        return a.T @ psi @ b + jnp.outer(q, g), psi
+
+    init = jnp.zeros_like(jnp.outer(q[-1], g[-1]))
+    _, psi = jax.lax.scan(impl, init, (a, b, q, g), reverse=True)
+    return psi
+
+
+@jax.jit
+def qsm_mul_psi_parallel(a, b, q, g):
+    B = jnp.einsum("nj,nk->njk", q, g)
+    _, _, psi = jax.lax.associative_scan(
+        _two_sided_affine_combine, (a.mT, b.mT, B), reverse=True
+    )
+    return _shift_bwd(psi)
+
+
+def _two_sided_affine_combine(left, right):
+    (al, bl, Bl), (ar, br, Br) = left, right
+    return ar @ al, br @ bl, ar @ Bl @ br.mT + Br
+
+
 @jax.jit
 def lower_matmul(p, q, a, x):
     def impl(f, data):