From b16ed0b9f798bec6df9414660492af52780d03f3 Mon Sep 17 00:00:00 2001 From: dgegen Date: Sun, 6 Sep 2026 16:38:07 +0200 Subject: [PATCH 1/2] Add memory-efficient mean/variance-only predict() path predict() routed through condition(), which always builds a dense N_test x N_test covariance (and its Cholesky factor) even when the caller only wanted .loc/diag(.variance). Add a fast path for return_var=True that reuses the existing training-set factorization and the kernel's evaluate_diag, without materializing that matrix. - Build the training/test cross-covariance block once and reuse it for both mean and variance, instead of evaluating the kernel twice. - return_cov=True and .sample() are unchanged, still on the dense condition() path. - Preserve return_var's documented precedence over return_cov: the fast path only defers to the dense condition() path when return_cov is requested without return_var. - Factor the duplicated X_test shape check into _check_test_shapes(). Tests: equivalence vs. condition() across a stationary, Quasisep, and pytree-input (LatentKernel) kernel; edge cases; a dedicated test for the return_var/return_cov precedence, added after an interim version of this refactor was found to invert it during review. --- src/tinygp/gp.py | 88 +++++++++++++++++++++++-------- tests/test_gp.py | 131 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 21 deletions(-) diff --git a/src/tinygp/gp.py b/src/tinygp/gp.py index 1feac47d..bd8e94b6 100644 --- a/src/tinygp/gp.py +++ b/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,50 @@ 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 + + # return_var=True: build the training/test cross-covariance block + # once (`Ks`) and reuse it for both the mean and the variance, rather + # than evaluating the kernel a second time via `_condition` -- this + # matters for kernels that are expensive to evaluate (e.g. ones that + # differentiate through another kernel). + alpha = self._get_alpha(y) + alpha = self.solver.solve_triangular(alpha, transpose=True) + + pred_kernel = self.kernel if kernel is None else kernel + if X_test is None: + Ks = pred_kernel(self.X, self.X) + Kss_diag = pred_kernel(self.X) + mean_offset = self.loc if include_mean else jnp.zeros_like(self.loc) + else: + 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 +414,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" + ) diff --git a/tests/test_gp.py b/tests/test_gp.py index 9268b6ad..3029a216 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -74,3 +74,134 @@ 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): + # Regression test: per the `predict` docstring, if `return_var` is True, + # `return_cov` must be ignored -- 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) From 771fb7e16d171c1a00e1434f74d4e6c9c13d9239 Mon Sep 17 00:00:00 2001 From: dgegen Date: Sun, 6 Sep 2026 17:55:04 +0200 Subject: [PATCH 2/2] predict(): avoid dense N_test x N_test work for return_var=True predict(y, X_test, return_var=True) routed through condition(), which always builds a dense N_test x N_test covariance (and, for DirectSolver, factorizes it) even when the caller only wants .loc/diag(.variance). - New test points (X_test given): build the training/test cross-covariance block once and reuse it for both the mean (alpha @ Ks) and the variance (Kss_diag - sum(solve_triangular(Ks)**2, axis=0)), instead of evaluating the kernel twice. - Predicting at the training points (X_test=None): reuse _condition()'s existing O(N) mean shortcut, and add Solver.condition_diag() so each solver can choose how to compute the variance. QuasisepSolver overrides it to reuse its quasiseparable structure (the same special case condition() already has), computing the diagonal in O(N) instead of building a dense N_train x N_train matrix -- this matters because Quasisep solvers exist specifically for large N. - return_cov=True and anything needing the full conditioned GaussianProcess (e.g. .sample()) are unchanged, still on the dense condition() path. - Preserve return_var's documented precedence over return_cov. - Factor the duplicated X_test shape check into _check_test_shapes(). Tests: equivalence vs. condition() across a stationary, Quasisep, and pytree-input (LatentKernel) kernel; edge cases; return_var/return_cov precedence; Quasisep prediction at training points at N large enough to distinguish the O(N) path from a dense fallback. --- src/tinygp/gp.py | 41 +++++++++++++++------------ src/tinygp/solvers/quasisep/solver.py | 22 ++++++++++++++ src/tinygp/solvers/solver.py | 31 ++++++++++++++++++++ tests/test_gp.py | 20 +++++++++++-- 4 files changed, 93 insertions(+), 21 deletions(-) diff --git a/src/tinygp/gp.py b/src/tinygp/gp.py index bd8e94b6..b00203b9 100644 --- a/src/tinygp/gp.py +++ b/src/tinygp/gp.py @@ -264,27 +264,32 @@ def predict( _, _, mean_value = self._condition(y, X_test, include_mean, kernel) return mean_value - # return_var=True: build the training/test cross-covariance block - # once (`Ks`) and reuse it for both the mean and the variance, rather - # than evaluating the kernel a second time via `_condition` -- this - # matters for kernels that are expensive to evaluate (e.g. ones that - # differentiate through another kernel). + 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) - pred_kernel = self.kernel if kernel is None else kernel - if X_test is None: - Ks = pred_kernel(self.X, self.X) - Kss_diag = pred_kernel(self.X) - mean_offset = self.loc if include_mean else jnp.zeros_like(self.loc) - else: - 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) - ) + 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) diff --git a/src/tinygp/solvers/quasisep/solver.py b/src/tinygp/solvers/quasisep/solver.py index c76f5ea7..bdb63137 100644 --- a/src/tinygp/solvers/quasisep/solver.py +++ b/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): diff --git a/src/tinygp/solvers/solver.py b/src/tinygp/solvers/solver.py index 68f927f0..9c3ad376 100644 --- a/src/tinygp/solvers/solver.py +++ b/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() diff --git a/tests/test_gp.py b/tests/test_gp.py index 3029a216..5e65d69e 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -191,9 +191,8 @@ def test_predict_edge_cases(random): def test_predict_return_var_takes_priority_over_return_cov(random): - # Regression test: per the `predict` docstring, if `return_var` is True, - # `return_cov` must be ignored -- the second return value must be the - # 1-D variance, not the 2-D covariance, even when both flags are set. + # 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) @@ -205,3 +204,18 @@ def test_predict_return_var_takes_priority_over_return_cov(random): 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)