Add memory-efficient fast path for predict(return_var=True)

src/tinygp/gp.py

@@ -173,22 +173,7 @@ def condition(
             the :class:`GaussianProcess` object describing the conditional
             distribution evaluated at ``X_test``.
         """
-        # If X_test is provided, we need to check that the tree structure
-        # matches that of the input data, and that the shapes are all compatible
-        # (i.e. the dimension of the inputs must match). This is slightly
-        # convoluted since we need to support arbitrary pytrees.
-        if X_test is not None:
-            matches = jax.tree_util.tree_map(
-                lambda a, b: jnp.ndim(a) == jnp.ndim(b)
-                and jnp.shape(a)[1:] == jnp.shape(b)[1:],
-                self.X,
-                X_test,
-            )
-            if not jax.tree_util.tree_reduce(lambda a, b: a and b, matches):
-                raise ValueError(
-                    "`X_test` must have the same tree structure as the input `X`, "
-                    "and all but the leading dimension must have matching sizes"
-                )
+        _check_test_shapes(self.X, X_test)
 
         alpha, log_prob, mean_value = self._condition(y, X_test, include_mean, kernel)
         if kernel is None:
@@ -263,12 +248,55 @@ def predict(
             returned with shape ``(N_test,)`` or ``(N_test, N_test)``
             respectively.
         """
-        _, cond = self.condition(y, X_test, kernel=kernel, include_mean=include_mean)
-        if return_var:
-            return cond.loc, cond.variance
-        if return_cov:
+        _check_test_shapes(self.X, X_test)
+
+        # `return_var` takes priority over `return_cov` (see docstring above),
+        # so only take the dense `condition()` path — which materializes the
+        # full N_test x N_test covariance — when the covariance itself was
+        # requested.
+        if return_cov and not return_var:
+            _, cond = self.condition(
+                y, X_test, kernel=kernel, include_mean=include_mean
+            )
             return cond.loc, cond.covariance
-        return cond.loc
+
+        if not return_var:
+            _, _, mean_value = self._condition(y, X_test, include_mean, kernel)
+            return mean_value
+
+        pred_kernel = self.kernel if kernel is None else kernel
+
+        if X_test is None:
+            # Predicting at the training points: reuse `_condition`'s O(N)
+            # mean shortcut, and let the solver pick how to compute the
+            # variance -- e.g. QuasisepSolver avoids a dense matrix here.
+            _, _, 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)
+            return mean_value, var_value
+
+        # Predicting at new test points: build `Ks` once and reuse it for
+        # both the mean and the variance, instead of evaluating the kernel
+        # twice via `_condition` -- matters for expensive kernels (e.g. ones
+        # that differentiate through another kernel).
+        alpha = self._get_alpha(y)
+        alpha = self.solver.solve_triangular(alpha, transpose=True)
+
+        Ks = pred_kernel(self.X, X_test)
+        Kss_diag = pred_kernel(X_test)
+        mean_offset = (
+            jax.vmap(self.mean_function)(X_test)
+            if include_mean
+            else jnp.zeros_like(Kss_diag)
+        )
+
+        mean_value = jnp.dot(alpha, Ks) + mean_offset
+        A = self.solver.solve_triangular(Ks)
+        var_value = (
+            Kss_diag - jnp.sum(jnp.square(A), axis=0) + _default_diag(mean_value)
+        )
+        return mean_value, var_value
 
     def sample(
         self,
@@ -391,3 +419,26 @@ def _default_diag(reference: JAXArray) -> JAXArray:
     give sensible results in general.
     """
     return jnp.sqrt(jnp.finfo(reference).eps)
+
+
+def _check_test_shapes(X: JAXArray, X_test: JAXArray | None) -> None:
+    """Check that ``X_test`` is compatible with the training inputs ``X``
+
+    If ``X_test`` is provided, we need to check that the tree structure
+    matches that of the input data, and that the shapes are all compatible
+    (i.e. the dimension of the inputs must match). This is slightly
+    convoluted since we need to support arbitrary pytrees.
+    """
+    if X_test is None:
+        return
+    matches = jax.tree_util.tree_map(
+        lambda a, b: jnp.ndim(a) == jnp.ndim(b)
+        and jnp.shape(a)[1:] == jnp.shape(b)[1:],
+        X,
+        X_test,
+    )
+    if not jax.tree_util.tree_reduce(lambda a, b: a and b, matches):
+        raise ValueError(
+            "`X_test` must have the same tree structure as the input `X`, "
+            "and all but the leading dimension must have matching sizes"
+        )

src/tinygp/solvers/quasisep/solver.py

@@ -138,6 +138,28 @@ def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> An
         A = self.solve_triangular(Ks)
         return Kss - A.transpose() @ A
 
+    def condition_diag(
+        self, kernel: Kernel, X_test: JAXArray | None, noise: Noise
+    ) -> JAXArray:
+        """The diagonal of the covariance matrix for a conditional GP
+
+        This reuses the same quasiseparable special case as :func:`condition`:
+        when predicting at the input coordinates with a
+        :class:`tinygp.kernels.quasisep.Quasisep` kernel, the diagonal can be
+        computed in ``O(N)`` (or ``O(N log N)`` with the parallel algorithms)
+        without ever materializing a dense matrix. Otherwise, this falls back on
+        :func:`tinygp.solvers.solver.Solver.condition_diag`.
+        """
+        from tinygp.kernels.quasisep import Quasisep
+
+        if X_test is None and isinstance(kernel, Quasisep):
+            M = kernel.to_symm_qsm(self.X)
+            delta = (self.factor.inv() @ M).gram()
+            M += noise.to_qsm()
+            return M.diag.d - delta.diag.d
+
+        return super().condition_diag(kernel, X_test, noise)
+
 
 def _check_sorted(X: JAXArray) -> None:
     if np.any(np.diff(X) < 0.0):

src/tinygp/solvers/solver.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 import equinox as eqx
+import jax.numpy as jnp
 
 from tinygp.helpers import JAXArray
 from tinygp.kernels.base import Kernel
@@ -80,3 +81,33 @@ def dot_triangular(self, y: JAXArray) -> JAXArray:
     @abstractmethod
     def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any:
         raise NotImplementedError
+
+    def condition_diag(
+        self, kernel: Kernel, X_test: JAXArray | None, noise: Noise
+    ) -> JAXArray:
+        """The diagonal of the covariance matrix for a conditional GP
+
+        Unlike :func:`condition`, which returns the full ``N_test x N_test``
+        conditional covariance matrix, this only computes its diagonal,
+        reusing this solver's existing factorization. The default
+        implementation below does this in ``O(N_train * N_test)`` time and
+        memory, without ever materializing a dense ``N_test x N_test``
+        matrix. Subclasses can override this to use a more efficient,
+        solver-specific algorithm where one is available (see e.g.
+        :class:`tinygp.solvers.QuasisepSolver`).
+
+        Args:
+            kernel: The kernel for the covariance between the observed and
+                predicted data.
+            X_test: The coordinates of the predicted points. Defaults to the
+                input coordinates.
+            noise: The noise model for the predicted process.
+        """
+        if X_test is None:
+            Ks = kernel(self.X, self.X)  # type: ignore
+            Kss_diag = kernel(self.X)  # type: ignore
+        else:
+            Ks = kernel(self.X, X_test)  # type: ignore
+            Kss_diag = kernel(X_test)
+        A = self.solve_triangular(Ks)
+        return Kss_diag - jnp.sum(jnp.square(A), axis=0) + noise.diagonal()

tests/test_gp.py

@@ -74,3 +74,148 @@ def distance(self, X1, X2):
     if tree:
         with pytest.raises(ValueError):
             gp.condition(y, {"x": X[0]})
+        with pytest.raises(ValueError):
+            gp.predict(y, {"x": X[0]})
+    else:
+        with pytest.raises(ValueError):
+            gp.predict(y, X[0])
+
+
+class LatentKernel(kernels.Kernel):
+    kernel: kernels.Kernel
+    coeff_prim: jax.Array
+    coeff_deriv: jax.Array
+
+    def __init__(self, kernel, coeff_prim, coeff_deriv):
+        self.kernel = kernel
+        self.coeff_prim, self.coeff_deriv = jnp.broadcast_arrays(
+            jnp.asarray(coeff_prim), jnp.asarray(coeff_deriv)
+        )
+
+    def evaluate(self, X1, X2):
+        t1, label1 = X1
+        t2, label2 = X2
+        Kp = jax.grad(self.kernel.evaluate, argnums=0)
+        Kpp = jax.grad(Kp, argnums=1)
+        K = self.kernel.evaluate(t1, t2)
+        d2K_dx1dx2 = Kpp(t1, t2)
+        dK_dx2 = jax.grad(self.kernel.evaluate, argnums=1)(t1, t2)
+        dK_dx1 = Kp(t1, t2)
+        a1 = self.coeff_prim[label1]
+        a2 = self.coeff_prim[label2]
+        b1 = self.coeff_deriv[label1]
+        b2 = self.coeff_deriv[label2]
+        return a1 * a2 * K + a1 * b2 * dK_dx2 + b1 * a2 * dK_dx1 + b1 * b2 * d2K_dx1dx2
+
+
+@pytest.mark.parametrize(
+    "kernel_type",
+    ["stationary", "quasisep", "pytree"],
+)
+def test_predict_equivalence(random, kernel_type):
+    with jax.enable_x64(True):
+        if kernel_type == "stationary":
+            X = jnp.sort(random.uniform(0, 10, (40, 2)), axis=0)
+            y = jnp.sin(X[:, 0]) + 0.1 * random.normal(size=len(X))
+            kernel = kernels.Matern32(1.5)
+            gp = GaussianProcess(kernel, X, diag=0.05, mean=jnp.sum)
+            X_test = jnp.sort(random.uniform(-2, 12, (60, 2)), axis=0)
+        elif kernel_type == "quasisep":
+            X = jnp.sort(random.uniform(0, 10, 50))
+            y = jnp.sin(X) + 0.1 * random.normal(size=len(X))
+            kernel = kernels.quasisep.SHO(omega=1.2, quality=2.5)
+            gp = GaussianProcess(kernel, X, diag=0.05)
+            X_test = jnp.sort(random.uniform(-2, 12, 70))
+        elif kernel_type == "pytree":
+            t_train = jnp.sort(random.uniform(0, 10, 30))
+            label_train = random.integers(0, 2, size=30)
+            X = (t_train, label_train)
+            y = jnp.sin(t_train) + 0.1 * random.normal(size=len(t_train))
+            base_k = kernels.ExpSquared(2.0)
+            kernel = LatentKernel(base_k, [1.0, 0.5], [-0.1, 0.3])
+            gp = GaussianProcess(kernel, X, diag=0.01)
+            t_test = jnp.sort(random.uniform(-1, 11, 45))
+            label_test = random.integers(0, 2, size=45)
+            X_test = (t_test, label_test)
+
+        _, cond = gp.condition(y, X_test)
+        mu_pred = gp.predict(y, X_test)
+        mu_var, var_pred = gp.predict(y, X_test, return_var=True)
+        mu_cov, cov_pred = gp.predict(y, X_test, return_cov=True)
+
+        assert_allclose(mu_pred, cond.loc)
+        assert_allclose(mu_var, cond.loc)
+        assert_allclose(mu_cov, cond.loc)
+        assert_allclose(var_pred, cond.variance)
+        assert_allclose(cov_pred, cond.covariance)
+
+
+def test_predict_edge_cases(random):
+    with jax.enable_x64(True):
+        X = jnp.sort(random.uniform(0, 10, 40))
+        y = jnp.sin(X)
+        kernel = kernels.Matern32(1.5)
+        gp = GaussianProcess(kernel, X, diag=0.05, mean=jnp.sin)
+
+        # X_test is None
+        _, cond_none = gp.condition(y)
+        mu_none, var_none = gp.predict(y, return_var=True)
+        assert_allclose(mu_none, cond_none.loc)
+        assert_allclose(var_none, cond_none.variance)
+
+        # X_test is exactly X
+        _, cond_x = gp.condition(y, X)
+        mu_x, var_x = gp.predict(y, X, return_var=True)
+        assert_allclose(mu_x, cond_x.loc)
+        assert_allclose(var_x, cond_x.variance)
+
+        # N_test = 1
+        X_1 = jnp.array([5.2])
+        _, cond_1 = gp.condition(y, X_1)
+        mu_1, var_1 = gp.predict(y, X_1, return_var=True)
+        assert_allclose(mu_1, cond_1.loc)
+        assert_allclose(var_1, cond_1.variance)
+
+        # include_mean=False
+        _, cond_nomean = gp.condition(y, X_1, include_mean=False)
+        mu_nomean, var_nomean = gp.predict(y, X_1, include_mean=False, return_var=True)
+        assert_allclose(mu_nomean, cond_nomean.loc)
+        assert_allclose(var_nomean, cond_nomean.variance)
+
+        # Custom cross kernel
+        cross_kernel = kernels.Exp(1.2)
+        _, cond_cross = gp.condition(y, X_1, kernel=cross_kernel)
+        mu_cross, var_cross = gp.predict(y, X_1, kernel=cross_kernel, return_var=True)
+        assert_allclose(mu_cross, cond_cross.loc)
+        assert_allclose(var_cross, cond_cross.variance)
+
+
+def test_predict_return_var_takes_priority_over_return_cov(random):
+    # The second return value must be the 1-D variance, not the 2-D
+    # covariance, even when both flags are set.
+    with jax.enable_x64(True):
+        X = jnp.sort(random.uniform(0, 10, 20))
+        y = jnp.sin(X)
+        gp = GaussianProcess(kernels.Matern32(1.5), X, diag=0.05)
+        X_test = jnp.sort(random.uniform(-1, 11, 6))
+
+        _, cond = gp.condition(y, X_test)
+        mu, out = gp.predict(y, X_test, return_var=True, return_cov=True)
+        assert out.shape == (6,)
+        assert_allclose(mu, cond.loc)
+        assert_allclose(out, cond.variance)
+
+
+def test_predict_quasisep_at_training_points(random):
+    # Must go through QuasisepSolver.condition_diag's quasiseparable
+    # shortcut, not fall back to a dense N x N cross-covariance matrix.
+    with jax.enable_x64(True):
+        X = jnp.sort(random.uniform(0, 10, 50))
+        y = jnp.sin(X) + 0.1 * random.normal(size=len(X))
+        kernel = kernels.quasisep.SHO(omega=1.2, quality=2.5)
+        gp = GaussianProcess(kernel, X, diag=0.05)
+
+        _, cond = gp.condition(y)
+        mu, var = gp.predict(y, return_var=True)
+        assert_allclose(mu, cond.loc)
+        assert_allclose(var, cond.variance)