From 984d0b21304f87e6b551acfd304fa9b3f3fa5bf1 Mon Sep 17 00:00:00 2001
From: Thomas Pinder <tompinder@live.co.uk>
Date: Fri, 7 Aug 2026 22:54:48 +0200
Subject: [PATCH] feat(state-space): dense joint predictive covariance for the
 smoothed posterior

Implements the acceptance criteria from #651: `StateSpacePrior.predict`,
`StateSpaceConjugateModel.predict`, and `StateSpacePosterior.__call__` now
support `covariance="dense"`, matching the dense `ConjugateModel` predictive
(mean and full covariance) to ~1e-8 for Matern-1/2, 3/2, 5/2.

- `rts_smoother` gains an opt-in `return_gains=True` that exposes its
  already-computed per-step smoother gains (no extra numerical work, just
  not discarding them). Existing 2-tuple callers are unaffected.
- New `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains
  those gains into the M x M joint covariance across test points, following
  the RTS smoother cross-covariance recursion (Sarkka & Solin 2019 SS12.2):
  Cov(x_i, x_j | y) = G_i...G_{j-1} P_j^smooth for i < j. It never inverts a
  gain product (ill-conditioned for widely separated points, since gains
  shrink with lag) and never forms an N x N gram over the training set --
  cost is O(N d^3) for the filter/smoother pass plus O(M^2 d^3) for the
  cross-covariance chaining, both linear in N.
- `StateSpacePrior.predict` (no conditioning data) returns the kernel's own
  dense gram for `covariance="dense"`, since the SDE is an exact
  representation of the kernel -- no Kalman machinery needed.

Design decision (the issue underspecifies this): `StateSpacePosterior.filtered`
/ `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep
raising NotImplementedError for `covariance="dense"`. Each filtered test
point conditions on a different information set (observations up to its own
timestamp), so a "joint" filtered covariance is not the dense
conjugate-predictive-shaped object the smoothed path now matches -- the
issue's acceptance criteria only compares against the dense conjugate
predictive, which is the smoothed quantity, and an existing repo test
(test_state_space_posterior_predict_filter_dense_raises) already pinned the
filtered-raises behaviour. Extending the filtered path is left to a future
issue if there's demand.

Tests: mean+covariance equivalence to the dense ConjugateModel (parametrized
over Matern12/32/52 and jitter), diagonal/dense consistency, caller-order
preservation, the M=1 degenerate case, a joint-sampling smoke test, jit and
grad cleanliness, a machine-precision numpy-oracle check on the exposed
smoother gains, and a larger-N robustness smoke test. Existing tests that
pinned the old "dense raises" behaviour for the now-supported paths are
updated to positive equivalence/plumbing checks instead.

Note for reviewers: gpjax/state_space/inference.py::rts_smoother is
implemented against its current covariance-form recursion, not a
square-root rework (that's issue #668, tracked separately). If #668 lands a
true square-root smoother, this cross-covariance recursion is worth
revisiting -- a square-root form may simplify the derivation -- but nothing
here blocks on it.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
---
 CHANGELOG.md                                |  25 ++
 gpjax/state_space/conditioning.py           |  38 +--
 gpjax/state_space/gps.py                    |  53 +++--
 gpjax/state_space/inference.py              |  36 ++-
 gpjax/state_space/prediction.py             | 224 ++++++++++++++++--
 tests/test_state_space/test_conditioning.py |  72 ++++--
 tests/test_state_space/test_gps.py          |  88 ++++++-
 tests/test_state_space/test_prediction.py   | 244 ++++++++++++++++++++
 tests/test_state_space/test_smoother.py     |  49 +++-
 9 files changed, 742 insertions(+), 87 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0c18588e7..dd3c84ad2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **Dense joint predictive covariance for state-space GPs.**
+  `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and
+  `StateSpacePosterior.__call__` (`gpjax.state_space`) now accept
+  `covariance="dense"`, returning the full joint covariance across test
+  points rather than marginal variances only
+  ([#651](https://github.com/thomaspinder/GPJax/issues/651)). For the
+  unconditioned prior this is just the kernel's own dense gram — the
+  state-space SDE is an exact representation with no training data to
+  marginalise out. For the conditioned (smoothed) posterior it is built from
+  the RTS smoother's cross-covariance recursion (Särkkä & Solin 2019 §12.2):
+  `rts_smoother` gains an opt-in `return_gains=True` that exposes its
+  already-computed per-step smoother gains, and a new
+  `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains them
+  into the `M x M` test-point covariance. This keeps the state-space
+  formulation's linear-in-`N` cost — no `N x N` gram over the training set is
+  ever formed — with the cross-covariance work landing at `O(M^2 d^3)`,
+  independent of `N` and no larger than the `O(M^2)` already required to
+  store the dense output. `StateSpacePosterior.filtered` /
+  `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep
+  raising `NotImplementedError` for `covariance="dense"`: each of their test
+  points conditions on a different information set (observations up to its
+  own timestamp), so a "joint" filtered covariance is not the dense
+  conjugate-predictive-shaped object the smoothed path now matches, and
+  extending it is left to a future issue if there is demand for it.
+
 - **`gpjax.fit_natgrads` and the `gpjax.natural_gradients` module.** Trains a variational
   family by alternating one natural-gradient step on the variational distribution with
   one step of a supplied Optax optimiser on everything else — kernel and likelihood
diff --git a/gpjax/state_space/conditioning.py b/gpjax/state_space/conditioning.py
index fc3ebb27b..83b98695c 100644
--- a/gpjax/state_space/conditioning.py
+++ b/gpjax/state_space/conditioning.py
@@ -66,11 +66,19 @@ class StateSpacePosterior(Posterior):
     gives the causal (filter-only) alternative, and
     :attr:`log_marginal_likelihood` gives the evidence from the same filter.
 
-    Predictive contract (v1): every query returns diagonal (marginal)
-    covariance only. The marginals are exact; a dense joint predictive is not
-    implemented in v1. This predictive is therefore not Liskov-substitutable
-    for a dense :class:`gpjax.conditioning.ExactPosterior` predictive, and
-    ``covariance="dense"`` raises :class:`NotImplementedError`.
+    Predictive contract: :meth:`__call__` (the smoothed predictive) supports
+    both ``covariance="diagonal"`` (marginal variances, the default) and
+    ``covariance="dense"`` (the full joint covariance across test points, via
+    the RTS smoother's cross-covariance recursion — see
+    :func:`gpjax.state_space.prediction._dense_smoothed_test_covariance`), so
+    it is Liskov-substitutable for a dense
+    :class:`gpjax.conditioning.ExactPosterior` predictive. :meth:`filtered`
+    (the causal predictive) has no dense joint form in v1 — each test point
+    conditions on a *different* information set (data up to its own
+    timestamp), so a "joint" filtered covariance is not the same kind of
+    object as the dense conjugate predictive it would otherwise be compared
+    against — and ``covariance="dense"`` there still raises
+    :class:`NotImplementedError`.
 
     Time ordering: the predictive queries merge and sort the train and test
     grids internally, so they are order-insensitive.
@@ -150,18 +158,18 @@ def __call__(
 
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
-            covariance: Must be ``"diagonal"``; the v1 state-space predictive
-                has no dense joint form.
+            covariance: ``"diagonal"`` returns the marginal variances only;
+                ``"dense"`` returns the full ``M x M`` joint covariance, built
+                from the RTS smoother's cross-covariance recursion (Särkkä &
+                Solin 2019 §12.2) rather than a dense ``N x N`` gram over the
+                training set.
 
         Returns:
             GaussianDistribution: The smoothed predictive, carrying an
-                ``lx.DiagonalLinearOperator`` scale.
-
-        Raises:
-            NotImplementedError: If ``covariance="dense"``.
+                ``lx.DiagonalLinearOperator`` scale when ``covariance`` is
+                ``"diagonal"`` or an ``lx.MatrixLinearOperator`` when it is
+                ``"dense"``.
         """
-        if covariance != "diagonal":
-            raise _dense_not_implemented("prediction")
         from gpjax.state_space.prediction import predict_smoothed
 
         return predict_smoothed(
@@ -169,6 +177,7 @@ def __call__(
             self.train_data,
             test_inputs,
             observation_mask=self.observation_mask,
+            covariance=covariance,
         )
 
     def predict(
@@ -187,7 +196,8 @@ def predict(
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
             train_data: Ignored.
-            covariance: Must be ``"diagonal"``.
+            covariance: ``"diagonal"`` for marginal variances or ``"dense"``
+                for the full joint covariance.
 
         Returns:
             GaussianDistribution: The smoothed predictive.
diff --git a/gpjax/state_space/gps.py b/gpjax/state_space/gps.py
index e6813c6d6..49b82d80a 100644
--- a/gpjax/state_space/gps.py
+++ b/gpjax/state_space/gps.py
@@ -20,6 +20,7 @@
 from gpjax.distributions import GaussianDistribution
 from gpjax.gps import ConjugateModel, Prior
 from gpjax.likelihoods import Gaussian, MultiOutputGaussian
+from gpjax.linalg.utils import add_jitter
 from gpjax.state_space.conditioning import StateSpacePosterior
 from gpjax.typing import Array
 
@@ -27,15 +28,14 @@
 class StateSpacePrior(Prior):
     """Prior for a state-space (Markovian) GP.
 
-    Identical to ``gpjax.gps.Prior`` except predictions are diagonal-only
-    (the prior is stationary in time, so off-diagonal covariance carries no
-    extra information for v1's diagonal-only predictive contract).
-
-    **Predictive contract (v1):** prediction returns diagonal (marginal)
-    covariance only; the marginals are exact. A dense joint predictive is not
-    implemented in v1 and is tracked as a follow-up. This predictive is
-    therefore not Liskov-substitutable for a dense
-    dense ``gpjax.gps.ConjugateModel`` predictive.
+    Predictions match ``gpjax.gps.Prior`` for both covariance modes.
+    ``covariance="diagonal"`` (the default) uses the SDE's stationary state
+    covariance directly rather than routing through the kernel, since the
+    prior is stationary in time; ``covariance="dense"`` returns the kernel's
+    own dense gram over the test inputs, since the state-space SDE is an
+    *exact* representation of the kernel with no training data to
+    marginalise out. Either way, this predictive is Liskov-substitutable for
+    the dense ``gpjax.gps.Prior`` predictive.
 
     Example:
         >>> import gpjax as gpx
@@ -52,13 +52,16 @@ def __call__(self, test_inputs, *, covariance="diagonal"):
         return self.predict(test_inputs, covariance=covariance)
 
     def predict(self, test_inputs, *, covariance="diagonal"):
-        if covariance != "diagonal":
-            raise NotImplementedError(
-                "State-space prior prediction returns diagonal (marginal) covariance "
-                "only; a dense joint predictive is not implemented in v1 and is "
-                "tracked as a follow-up. The marginal variances returned are exact, "
-                "so for diagonal-only use pass covariance='diagonal'."
+        mean_at_test = self.mean_function(test_inputs)
+        loc = jnp.atleast_1d(mean_at_test.squeeze())
+
+        if covariance == "dense":
+            gram_dense = add_jitter(
+                self.kernel.gram(test_inputs).as_matrix(), self.jitter
             )
+            scale = lx.MatrixLinearOperator(gram_dense)
+            return GaussianDistribution(loc=loc, scale=scale)
+
         from gpjax.state_space.kernels import to_sde
 
         sde = to_sde(self.kernel)
@@ -68,8 +71,6 @@ def predict(self, test_inputs, *, covariance="diagonal"):
         marginal_variance = (H @ P_inf @ H.T).squeeze() + self.jitter
 
         n_test = test_inputs.shape[0]
-        mean_at_test = self.mean_function(test_inputs)
-        loc = jnp.atleast_1d(mean_at_test.squeeze())
         scale = lx.DiagonalLinearOperator(jnp.full(n_test, marginal_variance))
         return GaussianDistribution(loc=loc, scale=scale)
 
@@ -94,11 +95,15 @@ class StateSpaceConjugateModel(ConjugateModel):
       - ``predict_filter``         : ``condition(D).filtered(t)``, the causal
         (filter-only) predictive
 
-    **Predictive contract (v1):** prediction returns diagonal (marginal)
-    covariance only; the marginals are exact. A dense joint predictive is not
-    implemented in v1 and is tracked as a follow-up. This predictive is
-    therefore not Liskov-substitutable for a dense
-    dense ``gpjax.gps.ConjugateModel`` predictive.
+    **Predictive contract:** ``predict``/``__call__`` (the smoothed
+    predictive) supports both ``covariance="diagonal"`` (marginal variances,
+    the default) and ``covariance="dense"`` (the full joint covariance across
+    test points, via the RTS smoother's cross-covariance recursion), so it is
+    Liskov-substitutable for the dense ``gpjax.gps.ConjugateModel``
+    predictive. ``predict_filter`` (the causal predictive) has no dense joint
+    form: each test point conditions on a different information set, so
+    ``covariance="dense"`` there still raises ``NotImplementedError`` — see
+    :class:`~gpjax.state_space.conditioning.StateSpacePosterior`.
 
     Example:
         >>> import gpjax as gpx
@@ -166,8 +171,8 @@ def predict(
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
             train_data: The observations to condition on.
-            covariance: Must be ``"diagonal"``; the v1 state-space predictive
-                has no dense joint form.
+            covariance: ``"diagonal"`` for marginal variances or ``"dense"``
+                for the full joint covariance.
             observation_mask: Optional boolean mask over the training points.
 
         Returns:
diff --git a/gpjax/state_space/inference.py b/gpjax/state_space/inference.py
index 29b5b9867..4b39b4d07 100644
--- a/gpjax/state_space/inference.py
+++ b/gpjax/state_space/inference.py
@@ -309,7 +309,7 @@ def skip_update(args):
     )
 
 
-def rts_smoother(sde, forward_outputs, time_steps):
+def rts_smoother(sde, forward_outputs, time_steps, *, return_gains: bool = False):
     r"""Square-root RTS smoother.
 
     Runs the standard Särkkä & Solin (2019) §10.7 backward recursion on the
@@ -328,6 +328,17 @@ def rts_smoother(sde, forward_outputs, time_steps):
     The last step has no future, so its smoothed state equals its filtered
     state.
 
+    The per-step gains :math:`G_i` are already computed inside the backward
+    scan to produce :math:`P^{\text{smooth}}_i`; ``return_gains`` just adds
+    them to the scan's output rather than discarding them, at no extra
+    numerical cost. They are the building block for the smoother
+    cross-covariance recursion
+    (:math:`\mathrm{Cov}(x_i, x_j \mid y_{1:T}) = G_i \cdots G_{j-1}
+    P^{\text{smooth}}_j` for :math:`i < j`, Särkkä & Solin 2019 §12.2) used by
+    :func:`gpjax.state_space.prediction._dense_test_covariance` to build the
+    joint predictive covariance without materialising an
+    :math:`N \times N` gram.
+
     Args:
         sde (LinearSDE): State-space SDE used in the forward pass;
             ``sde.discretise(dt)`` is called once per backward step.
@@ -337,11 +348,18 @@ def rts_smoother(sde, forward_outputs, time_steps):
         time_steps (Float[Array, "num_train"]): Same ``time_steps`` that drove
             the forward pass; ``time_steps[i+1]`` is the inter-step ``dt``
             between filtered index ``i`` and predicted index ``i + 1``.
+        return_gains (bool, keyword-only): If ``True``, also return the
+            per-step smoother gains :math:`G_i` for :math:`i = 0, \ldots,
+            \text{num\_train} - 2`.
 
     Returns:
         tuple: ``smoothed_means`` of shape ``Float[Array, "num_train state_dim"]``
             and ``smoothed_Ls`` of shape
-            ``Float[Array, "num_train state_dim state_dim"]``.
+            ``Float[Array, "num_train state_dim state_dim"]``. If
+            ``return_gains`` is ``True``, a third element ``smoother_gains`` of
+            shape ``Float[Array, "num_train-1 state_dim state_dim"]`` is
+            appended, with ``smoother_gains[i]`` the gain :math:`G_i`
+            connecting smoothed index ``i`` to smoothed index ``i + 1``.
 
     See plans/2026-04-21-state-space-gps-design.md §Stage 3.
     """
@@ -388,7 +406,7 @@ def backward_step(carry, scan_input):
         P_smoothed = 0.5 * (P_smoothed + P_smoothed.T)
         L_smoothed = _psd_sqrt(P_smoothed)
 
-        return (mean_smoothed, L_smoothed), (mean_smoothed, L_smoothed)
+        return (mean_smoothed, L_smoothed), (mean_smoothed, L_smoothed, smoother_gain)
 
     # Initial smoother carry = filtered state at the last step (no future).
     init_carry = (means_updated[-1], Ls_updated[-1])
@@ -408,7 +426,9 @@ def backward_step(carry, scan_input):
     _, smoothed_outputs_reversed = jax.lax.scan(
         backward_step, init_carry, backward_inputs_reversed
     )
-    smoothed_means_prefix, smoothed_Ls_prefix = smoothed_outputs_reversed
+    smoothed_means_prefix, smoothed_Ls_prefix, smoother_gains_prefix_reversed = (
+        smoothed_outputs_reversed
+    )
 
     # Un-reverse and append the last step (which equals the filtered state).
     smoothed_means = jnp.concatenate(
@@ -417,4 +437,10 @@ def backward_step(carry, scan_input):
     smoothed_Ls = jnp.concatenate(
         [jnp.flip(smoothed_Ls_prefix, axis=0), Ls_updated[-1:]], axis=0
     )
-    return smoothed_means, smoothed_Ls
+    if not return_gains:
+        return smoothed_means, smoothed_Ls
+
+    # smoother_gains[i] = G_i, connecting smoothed index i to i + 1; indices
+    # 0 .. num_train - 2, matching backward_inputs before the reverse.
+    smoother_gains = jnp.flip(smoother_gains_prefix_reversed, axis=0)
+    return smoothed_means, smoothed_Ls, smoother_gains
diff --git a/gpjax/state_space/prediction.py b/gpjax/state_space/prediction.py
index 8d7658da3..efdffe0e9 100644
--- a/gpjax/state_space/prediction.py
+++ b/gpjax/state_space/prediction.py
@@ -5,12 +5,15 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 import jax
 import jax.numpy as jnp
 import lineax as lx
 import paramax
 
 from gpjax.distributions import GaussianDistribution
+from gpjax.linalg.utils import add_jitter
 from gpjax.state_space.inference import _sqrt_filter_forward, rts_smoother
 from gpjax.state_space.kernels import to_sde
 
@@ -52,14 +55,190 @@ def _merge_grids(train_times, test_times, centred_targets, observation_mask):
     return sorted_times, sorted_targets, sorted_is_observed, sorted_is_test, merge_perm
 
 
-def predict_smoothed(posterior, train_data, test_inputs, *, observation_mask=None):
+def _test_segment_operators(smoother_gains, sorted_is_test, num_test):
+    r"""Cumulative smoother-gain products between consecutive test points.
+
+    The RTS smoother cross-covariance recursion (Särkkä & Solin 2019 §12.2)
+    gives :math:`\mathrm{Cov}(x_i, x_j \mid y_{1:T}) = G_i G_{i+1} \cdots
+    G_{j-1} P_j^{\text{smooth}}` for merged-grid indices :math:`i < j`. This
+    function computes, for each pair of *consecutive* test points in the
+    merged grid, the cumulative gain product :math:`G_{p_a} \cdots
+    G_{p_{a+1}-1}` connecting them — the reusable building block that
+    :func:`_dense_test_covariance` chains together for arbitrary (not just
+    adjacent) test-point pairs.
+
+    A single forward pass over the merged grid accumulates the running
+    product since the last test point was seen, resetting to the identity at
+    each test point and emitting the accumulated product the next time a test
+    point is reached. This costs :math:`O((N + M) d^3)` — linear in the
+    training set size :math:`N`, not quadratic — since it never touches a
+    pair of training points.
+
+    Args:
+        smoother_gains (Float[Array, "grid_size-1 state_dim state_dim"]): Per-step
+            smoother gains from :func:`gpjax.state_space.inference.rts_smoother`
+            with ``return_gains=True``, on the merged train-plus-test grid.
+        sorted_is_test (Bool[Array, " grid_size"]): Merged-grid mask, ``True``
+            at test positions.
+        num_test (int): Number of test points ``M`` (static; the size of the
+            ``True`` entries in ``sorted_is_test``).
+
+    Returns:
+        Float[Array, "num_test-1 state_dim state_dim"]: ``segments[a]`` is the
+            operator connecting the ``a``-th and ``(a + 1)``-th test point in
+            merged-grid order, so that
+            :math:`\mathrm{Cov}(x_{p_a}, x_{p_{a+1}}) = \text{segments}[a]\,
+            P_{p_{a+1}}^{\text{smooth}}`.
+    """
+    state_dim = smoother_gains.shape[-1]
+    is_test_bool = sorted_is_test.astype(bool)
+    identity = jnp.eye(state_dim)
+
+    def step(carry, step_inputs):
+        running_product, has_started = carry
+        gain, is_test_here, is_test_next = step_inputs
+        # Reset the running product to the identity right at a test point, so
+        # it starts fresh with `gain` as its first factor.
+        base = jnp.where(is_test_here, identity, running_product)
+        running_product_next = base @ gain
+        has_started_next = has_started | is_test_here
+        # Emit only once accumulation has actually started (skips any
+        # training-only prefix before the first test point).
+        emit = is_test_next & has_started_next
+        return (running_product_next, has_started_next), (running_product_next, emit)
+
+    init_carry = (identity, jnp.array(False))
+    step_inputs = (smoother_gains, is_test_bool[:-1], is_test_bool[1:])
+    _, (candidates, emit_mask) = jax.lax.scan(step, init_carry, step_inputs)
+
+    num_segments = num_test - 1
+    emit_indices = jnp.nonzero(emit_mask, size=num_segments)[0]
+    return candidates[emit_indices]
+
+
+def _dense_test_covariance(segments, state_covs_at_test, observation_matrix):
+    r"""Dense observation-space covariance across test points from smoother segments.
+
+    Builds the full :math:`M \times M` covariance in the latent *state* space
+    by chaining :func:`_test_segment_operators`' consecutive-pair operators —
+    without ever inverting a gain product, which would be ill-conditioned for
+    widely separated points since gains shrink towards zero with lag — then
+    projects through the observation matrix. For each test point ``b``, a
+    backward scan over the segments before it accumulates
+    :math:`\mathrm{Cov}(x_a, x_b)` for every ``a < b`` in one pass, mirroring
+    the RTS backward recursion itself but over the (typically far smaller)
+    test-point chain rather than the full merged grid. Cost is
+    :math:`O(M^2 d^3)`, independent of the training set size and no larger
+    than the :math:`O(M^2)` already required to store the dense output.
+
+    Args:
+        segments (Float[Array, "num_test-1 state_dim state_dim"]): Consecutive-pair
+            operators from :func:`_test_segment_operators`.
+        state_covs_at_test (Float[Array, "num_test state_dim state_dim"]): Smoothed
+            state covariances :math:`P_b^{\text{smooth}}` at each test point,
+            in the same (merged-grid) order as ``segments``.
+        observation_matrix (Float[Array, "1 state_dim"]): The SDE's observation
+            matrix :math:`H`.
+
+    Returns:
+        Float[Array, "num_test num_test"]: The dense observation-space
+            covariance matrix, in the same order as ``state_covs_at_test``.
+    """
+    num_test, state_dim, _ = state_covs_at_test.shape
+    H_row = observation_matrix.reshape(-1)
+    diag_variances = jnp.einsum("i,nij,j->n", H_row, state_covs_at_test, H_row)
+
+    if num_test <= 1:
+        # `segments` is empty (shape (0, state_dim, state_dim)); indexing it
+        # even abstractly (as `column_for` would, under `jax.lax.scan`
+        # tracing) is a static shape error in JAX, so skip it entirely.
+        return jnp.diag(diag_variances)
+
+    def column_for(b):
+        def body(carry, a_rev):
+            apply_segment = a_rev < b
+            new_carry = jnp.where(apply_segment, segments[a_rev] @ carry, carry)
+            return new_carry, new_carry
+
+        a_rev_range = jnp.arange(num_test - 2, -1, -1)
+        _, outputs = jax.lax.scan(body, state_covs_at_test[b], a_rev_range)
+        column = jnp.zeros((num_test, state_dim, state_dim))
+        return column.at[a_rev_range].set(outputs)
+
+    # columns[b, a] = Cov(x_a, x_b) for a < b; entries with a >= b are unused.
+    columns = jax.vmap(column_for)(jnp.arange(num_test))
+    obs_cross = jnp.einsum("i,baij,j->ba", H_row, columns, H_row)
+
+    lower_valid = jnp.tril(obs_cross, k=-1)
+    off_diagonal = lower_valid + lower_valid.T
+    return off_diagonal + jnp.diag(diag_variances)
+
+
+def _dense_smoothed_test_covariance(
+    smoother_gains,
+    sorted_is_test,
+    smoothed_Ls,
+    test_positions_in_sorted,
+    observation_matrix,
+):
+    r"""The dense test-test predictive covariance, in caller order.
+
+    Combines :func:`_test_segment_operators` and :func:`_dense_test_covariance`,
+    handling the sort into merged-grid order (both helpers above assume it)
+    and the sort back into caller order.
+
+    Args:
+        smoother_gains (Float[Array, "grid_size-1 state_dim state_dim"]): Per-step
+            smoother gains, as returned by ``rts_smoother(..., return_gains=True)``.
+        sorted_is_test (Bool[Array, " grid_size"]): Merged-grid mask, ``True``
+            at test positions.
+        smoothed_Ls (Float[Array, "grid_size state_dim state_dim"]): Smoothed
+            covariance square roots at every merged-grid position.
+        test_positions_in_sorted (Int[Array, " num_test"]): For caller-order
+            test index ``m``, the merged-grid position of that test point (as
+            returned by the ``inv_perm`` construction in :func:`predict_smoothed`).
+        observation_matrix (Float[Array, "1 state_dim"]): The SDE's observation
+            matrix :math:`H`.
+
+    Returns:
+        Float[Array, "num_test num_test"]: The dense observation-space
+            covariance matrix, in caller order.
+    """
+    num_test = test_positions_in_sorted.shape[0]
+    sort_order = jnp.argsort(test_positions_in_sorted)
+    grid_order_positions = test_positions_in_sorted[sort_order]
+    unsort_order = jnp.argsort(sort_order)
+
+    segments = _test_segment_operators(smoother_gains, sorted_is_test, num_test)
+    state_covs_grid_order = jax.vmap(lambda L: L @ L.T)(
+        smoothed_Ls[grid_order_positions]
+    )
+    covariance_grid_order = _dense_test_covariance(
+        segments, state_covs_grid_order, observation_matrix
+    )
+    return covariance_grid_order[unsort_order][:, unsort_order]
+
+
+def predict_smoothed(
+    posterior,
+    train_data,
+    test_inputs,
+    *,
+    observation_mask=None,
+    covariance: Literal["dense", "diagonal"] = "diagonal",
+):
     """Smoothed-latent prediction for state-space GPs.
 
     Returns a ``GaussianDistribution`` over the M test points whose mean is
-    the RTS-smoothed posterior mean and whose scale is a
-    ``lx.DiagonalLinearOperator`` carrying the smoothed marginal variance plus
-    ``prior.jitter``. Test inputs are returned in caller order regardless of
-    sorting.
+    the RTS-smoothed posterior mean. When ``covariance="diagonal"`` (the
+    default) the scale is a ``lx.DiagonalLinearOperator`` carrying the
+    smoothed marginal variance plus ``prior.jitter``. When
+    ``covariance="dense"`` the scale is an ``lx.MatrixLinearOperator``
+    carrying the full :math:`M \\times M` joint covariance, built from the
+    RTS smoother's cross-covariance recursion
+    (:func:`_dense_smoothed_test_covariance`) rather than a dense
+    :math:`N \\times N` gram over the training set. Test inputs are returned
+    in caller order regardless of sorting.
 
     See plans/2026-04-21-state-space-gps-design.md §Stage 4.
     """
@@ -88,7 +267,7 @@ def predict_smoothed(posterior, train_data, test_inputs, *, observation_mask=Non
         sorted_times,
         sorted_targets,
         sorted_is_observed,
-        _sorted_is_test,
+        sorted_is_test,
         merge_perm,
     ) = _merge_grids(train_times, test_times, centred_targets, observation_mask)
     # Time-step deltas on the merged grid.
@@ -101,7 +280,14 @@ def predict_smoothed(posterior, train_data, test_inputs, *, observation_mask=Non
         sorted_is_observed,
         sigma_eff,
     )
-    smoothed_means, smoothed_Ls = rts_smoother(sde, forward_outputs, sorted_time_steps)
+    if covariance == "dense":
+        smoothed_means, smoothed_Ls, smoother_gains = rts_smoother(
+            sde, forward_outputs, sorted_time_steps, return_gains=True
+        )
+    else:
+        smoothed_means, smoothed_Ls = rts_smoother(
+            sde, forward_outputs, sorted_time_steps
+        )
 
     # Recover test-point positions in caller order via the inverse permutation.
     inv_perm = jnp.argsort(merge_perm)
@@ -109,20 +295,32 @@ def predict_smoothed(posterior, train_data, test_inputs, *, observation_mask=Non
 
     H = sde.observation_matrix
     test_smoothed_means = smoothed_means[test_positions_in_sorted]
-    test_smoothed_Ls = smoothed_Ls[test_positions_in_sorted]
     test_observation_means = jnp.einsum("ij,mj->mi", H, test_smoothed_means).squeeze(-1)
-    test_observation_variances = jax.vmap(lambda L: (H @ (L @ L.T) @ H.T).squeeze())(
-        test_smoothed_Ls
-    )
 
     # Re-add mean function at test points + add jitter.
     mean_at_test = prior.mean_function(test_inputs).squeeze(-1)
     test_predicted_means = test_observation_means + mean_at_test
-    test_predicted_variances = test_observation_variances + prior.jitter
+
+    if covariance == "dense":
+        test_covariance = _dense_smoothed_test_covariance(
+            smoother_gains,
+            sorted_is_test,
+            smoothed_Ls,
+            test_positions_in_sorted,
+            H,
+        )
+        scale = lx.MatrixLinearOperator(add_jitter(test_covariance, prior.jitter))
+    else:
+        test_smoothed_Ls = smoothed_Ls[test_positions_in_sorted]
+        test_observation_variances = jax.vmap(
+            lambda L: (H @ (L @ L.T) @ H.T).squeeze()
+        )(test_smoothed_Ls)
+        test_predicted_variances = test_observation_variances + prior.jitter
+        scale = lx.DiagonalLinearOperator(test_predicted_variances)
 
     return GaussianDistribution(
         loc=test_predicted_means,
-        scale=lx.DiagonalLinearOperator(test_predicted_variances),
+        scale=scale,
     )
 
 
diff --git a/tests/test_state_space/test_conditioning.py b/tests/test_state_space/test_conditioning.py
index a523f4d00..6c8260ee4 100644
--- a/tests/test_state_space/test_conditioning.py
+++ b/tests/test_state_space/test_conditioning.py
@@ -135,6 +135,29 @@ def test_conditioned_call_matches_predict_smoothed(kernel_class):
     )
 
 
+@pytest.mark.parametrize(
+    "kernel_class",
+    [gpx.kernels.Matern12, gpx.kernels.Matern32, gpx.kernels.Matern52],
+)
+def test_conditioned_call_dense_matches_predict_smoothed(kernel_class):
+    """Same plumbing identity as above, for the dense joint covariance."""
+    model = _build_model(kernel_class)
+    train_data = _build_train_data()
+
+    conditioned = model.condition(train_data)(_TEST_INPUTS, covariance="dense")
+    reference = predict_smoothed(model, train_data, _TEST_INPUTS, covariance="dense")
+
+    assert isinstance(conditioned.scale, lx.MatrixLinearOperator)
+    np.testing.assert_allclose(
+        np.asarray(conditioned.mean), np.asarray(reference.mean), atol=1e-12
+    )
+    np.testing.assert_allclose(
+        np.asarray(conditioned.covariance_matrix),
+        np.asarray(reference.covariance_matrix),
+        atol=1e-12,
+    )
+
+
 def test_conditioned_call_defaults_to_diagonal():
     """``covariance`` defaults to diagonal; the v1 contract has no dense form."""
     model = _build_model()
@@ -210,44 +233,59 @@ def test_condition_threads_the_observation_mask():
 
 
 # ---------------------------------------------------------------------------
-# The dense rejection
+# The dense joint predictive (smoothed), and the still-rejected filtered one
 # ---------------------------------------------------------------------------
 
 
-def test_conditioned_call_dense_raises_with_actionable_message():
+def test_conditioned_call_dense_returns_matrix_operator():
     model = _build_model()
     train_data = _build_train_data()
     conditioned = model.condition(train_data)
 
-    with pytest.raises(NotImplementedError) as excinfo:
-        conditioned(_TEST_INPUTS, covariance="dense")
+    dist = conditioned(_TEST_INPUTS, covariance="dense")
 
-    message = str(excinfo.value)
-    assert "diagonal" in message
-    assert "not implemented in v1" in message
-    assert "covariance='diagonal'" in message
+    assert isinstance(dist, GaussianDistribution)
+    assert isinstance(dist.scale, lx.MatrixLinearOperator)
+    assert dist.covariance_matrix.shape == (
+        _TEST_INPUTS.shape[0],
+        _TEST_INPUTS.shape[0],
+    )
+    np.testing.assert_allclose(
+        np.asarray(dist.covariance_matrix),
+        np.asarray(dist.covariance_matrix).T,
+        atol=1e-10,
+    )
 
 
-def test_conditioned_filtered_dense_raises_with_actionable_message():
+def test_conditioned_predict_dense_matches_call():
+    """``predict`` is sugar for ``__call__``; the dense mode is no exception."""
     model = _build_model()
     train_data = _build_train_data()
     conditioned = model.condition(train_data)
 
-    with pytest.raises(NotImplementedError) as excinfo:
-        conditioned.filtered(_TEST_INPUTS, covariance="dense")
+    sugar = conditioned.predict(_TEST_INPUTS, covariance="dense")
+    explicit = conditioned(_TEST_INPUTS, covariance="dense")
 
-    message = str(excinfo.value)
-    assert "predict_filter" in message
-    assert "covariance='diagonal'" in message
+    np.testing.assert_array_equal(
+        np.asarray(sugar.covariance_matrix), np.asarray(explicit.covariance_matrix)
+    )
 
 
-def test_conditioned_predict_dense_raises():
+def test_conditioned_filtered_dense_raises_with_actionable_message():
+    """Unlike the smoothed predictive, ``filtered`` has no dense joint form:
+    each test point conditions on a different information set, so a "joint"
+    filtered covariance is not comparable to the dense conjugate predictive
+    the smoothed path now matches. See ``StateSpacePosterior``'s docstring."""
     model = _build_model()
     train_data = _build_train_data()
     conditioned = model.condition(train_data)
 
-    with pytest.raises(NotImplementedError, match=r"diagonal"):
-        conditioned.predict(_TEST_INPUTS, covariance="dense")
+    with pytest.raises(NotImplementedError) as excinfo:
+        conditioned.filtered(_TEST_INPUTS, covariance="dense")
+
+    message = str(excinfo.value)
+    assert "predict_filter" in message
+    assert "covariance='diagonal'" in message
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/test_state_space/test_gps.py b/tests/test_state_space/test_gps.py
index 0d61200bd..3f5334d3a 100644
--- a/tests/test_state_space/test_gps.py
+++ b/tests/test_state_space/test_gps.py
@@ -26,14 +26,45 @@ def test_state_space_prior_predict_diagonal_returns_gaussian_distribution():
     np.testing.assert_allclose(np.asarray(dist.mean), 0.0, atol=1e-12)
 
 
-def test_state_space_prior_predict_dense_raises():
+@pytest.mark.parametrize(
+    "kernel_class",
+    [gpx.kernels.Matern12, gpx.kernels.Matern32, gpx.kernels.Matern52],
+)
+def test_state_space_prior_predict_dense_matches_kernel_gram(kernel_class):
+    """The prior has no data to condition on, so the dense joint is exactly
+    the kernel's own gram — the SDE is an exact representation, not an
+    approximation."""
+    lengthscale, variance = 1.2, 0.9
     prior = StateSpacePrior(
         mean_function=gpx.mean_functions.Zero(),
-        kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
+        kernel=kernel_class(lengthscale=lengthscale, variance=variance),
+        jitter=1e-8,
+    )
+    Xtest = jnp.array([[0.0], [1.0], [3.5], [3.5], [10.0]])
+    dist = prior.predict(Xtest, covariance="dense")
+    assert isinstance(dist, GaussianDistribution)
+    assert isinstance(dist.scale, lx.MatrixLinearOperator)
+
+    dense_kernel = kernel_class(lengthscale=lengthscale, variance=variance)
+    expected_cov = dense_kernel.gram(Xtest).as_matrix() + 1e-8 * jnp.eye(5)
+    np.testing.assert_allclose(
+        np.asarray(dist.covariance_matrix), np.asarray(expected_cov), atol=1e-10
+    )
+    np.testing.assert_allclose(np.asarray(dist.mean), 0.0, atol=1e-12)
+
+
+def test_state_space_prior_predict_dense_matches_diagonal_on_the_diagonal():
+    prior = StateSpacePrior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=2.0),
+        jitter=1e-6,
+    )
+    Xtest = jnp.linspace(0.0, 5.0, 7).reshape(-1, 1)
+    dense_dist = prior.predict(Xtest, covariance="dense")
+    diagonal_dist = prior.predict(Xtest, covariance="diagonal")
+    np.testing.assert_allclose(
+        np.asarray(dense_dist.variance), np.asarray(diagonal_dist.variance), atol=1e-12
     )
-    Xtest = jnp.array([[0.0], [1.0]])
-    with pytest.raises(NotImplementedError, match=r"diagonal|dense"):
-        prior.predict(Xtest, covariance="dense")
 
 
 def test_state_space_prior_jitter_is_added_to_marginal_variance():
@@ -73,7 +104,34 @@ def test_state_space_conjugate_posterior_construction():
     assert posterior.likelihood is likelihood
 
 
-def test_state_space_conjugate_posterior_predict_dense_raises():
+def test_state_space_conjugate_posterior_predict_dense_matches_condition_call():
+    """``predict`` with ``covariance="dense"`` is sugar over ``condition(D)(t)``."""
+    prior = StateSpacePrior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=gpx.kernels.Matern12(lengthscale=1.0, variance=1.0),
+        jitter=1e-6,
+    )
+    likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1)
+    posterior = StateSpaceConjugateModel(prior=prior, likelihood=likelihood)
+
+    train_X = jnp.linspace(0.0, 1.0, 5).reshape(-1, 1)
+    train_y = jnp.zeros((5, 1))
+    train_data = gpx.Dataset(X=train_X, y=train_y)
+    Xtest = jnp.array([[0.5], [0.8]])
+
+    sugar = posterior.predict(Xtest, train_data, covariance="dense")
+    explicit = posterior.condition(train_data)(Xtest, covariance="dense")
+
+    assert isinstance(sugar.scale, lx.MatrixLinearOperator)
+    np.testing.assert_array_equal(
+        np.asarray(sugar.covariance_matrix), np.asarray(explicit.covariance_matrix)
+    )
+    np.testing.assert_array_equal(np.asarray(sugar.mean), np.asarray(explicit.mean))
+
+
+def test_state_space_conjugate_posterior_predict_filter_dense_still_raises():
+    """The causal predictive keeps rejecting ``covariance="dense"`` — see
+    ``StateSpacePosterior.filtered``'s docstring for why."""
     prior = StateSpacePrior(
         mean_function=gpx.mean_functions.Zero(),
         kernel=gpx.kernels.Matern12(lengthscale=1.0, variance=1.0),
@@ -86,7 +144,7 @@ def test_state_space_conjugate_posterior_predict_dense_raises():
     train_data = gpx.Dataset(X=train_X, y=train_y)
     Xtest = jnp.array([[0.5]])
     with pytest.raises(NotImplementedError, match=r"diagonal|dense"):
-        posterior.predict(Xtest, train_data, covariance="dense")
+        posterior.predict_filter(Xtest, train_data, covariance="dense")
 
 
 def test_state_space_prior_times_gaussian_returns_state_space_posterior():
@@ -130,7 +188,7 @@ def test_state_space_prior_times_gaussian_with_array_obs_stddev_raises():
         prior * likelihood
 
 
-def test_state_space_predict_rejects_dense_with_actionable_message():
+def test_state_space_predict_dense_returns_finite_psd_covariance():
     data = gpx.Dataset(
         X=jnp.linspace(0, 5, 10).reshape(-1, 1),
         y=jnp.sin(jnp.linspace(0, 5, 10)).reshape(-1, 1),
@@ -138,9 +196,15 @@ def test_state_space_predict_rejects_dense_with_actionable_message():
     posterior = StateSpacePrior(
         mean_function=gpx.mean_functions.Zero(),
         kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
+        jitter=1e-6,
     ) * gpx.likelihoods.Gaussian(obs_stddev=0.1)
 
-    with pytest.raises(NotImplementedError, match=r"diagonal-only|follow-up|v1"):
-        posterior.predict(
-            jnp.linspace(0, 5, 4).reshape(-1, 1), data, covariance="dense"
-        )
+    Xtest = jnp.linspace(0, 5, 4).reshape(-1, 1)
+    dist = posterior.predict(Xtest, data, covariance="dense")
+
+    cov = np.asarray(dist.covariance_matrix)
+    assert cov.shape == (4, 4)
+    assert np.all(np.isfinite(cov))
+    np.testing.assert_allclose(cov, cov.T, atol=1e-10)
+    eigenvalues = np.linalg.eigvalsh(cov)
+    assert np.all(eigenvalues > 0.0)
diff --git a/tests/test_state_space/test_prediction.py b/tests/test_state_space/test_prediction.py
index 3d3236fcd..462a60a5b 100644
--- a/tests/test_state_space/test_prediction.py
+++ b/tests/test_state_space/test_prediction.py
@@ -3,8 +3,11 @@
 import gpjax as gpx
 from gpjax.distributions import GaussianDistribution
 from gpjax.state_space.gps import StateSpacePrior
+from gpjax.state_space.inference import rts_smoother
 from gpjax.state_space.prediction import _merge_grids
+import jax
 import jax.numpy as jnp
+import jax.random as jr
 import lineax as lx
 import numpy as np
 import pytest
@@ -143,6 +146,247 @@ def test_state_space_posterior_predict_smoothed_matches_dense_gp(kernel_class, j
     np.testing.assert_allclose(ss_variances, dense_variances, atol=1e-5, rtol=1e-6)
 
 
+# ---------------------------------------------------------------------------
+# Dense joint predictive covariance (issue #651)
+# ---------------------------------------------------------------------------
+
+
+def _build_smoothed_dense_comparison(kernel_class, jitter, n_train=30, n_test=10):
+    """Shared setup for the dense-vs-dense-GP comparison tests below."""
+    lengthscale, variance, obs_stddev = 1.5, 0.8, 0.2
+    X_train, y_train = _build_matern12_dataset(
+        n=n_train,
+        lengthscale=lengthscale,
+        variance=variance,
+        obs_stddev=obs_stddev,
+    )
+    Xtest = jnp.linspace(-0.5, 10.5, n_test).reshape(-1, 1)
+    train_data = gpx.Dataset(X=X_train.reshape(-1, 1), y=y_train.reshape(-1, 1))
+
+    ss_posterior = StateSpacePrior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=kernel_class(lengthscale=lengthscale, variance=variance),
+        jitter=jitter,
+    ) * gpx.likelihoods.Gaussian(obs_stddev=obs_stddev)
+    dense_posterior = gpx.gps.Prior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=kernel_class(lengthscale=lengthscale, variance=variance),
+        jitter=jitter,
+    ) * gpx.likelihoods.Gaussian(obs_stddev=obs_stddev)
+    return ss_posterior, dense_posterior, train_data, Xtest
+
+
+@pytest.mark.parametrize(
+    "kernel_class",
+    [gpx.kernels.Matern12, gpx.kernels.Matern32, gpx.kernels.Matern52],
+)
+@pytest.mark.parametrize("jitter", [0.0, 1e-6])
+def test_state_space_posterior_predict_smoothed_dense_matches_dense_gp_joint(
+    kernel_class, jitter
+):
+    """Acceptance criterion (issue #651): the dense joint predictive — mean
+    *and* full covariance — must match the dense ``ConjugateModel`` predictive
+    to ~1e-8."""
+    ss_posterior, dense_posterior, train_data, Xtest = _build_smoothed_dense_comparison(
+        kernel_class, jitter
+    )
+
+    ss_dist = ss_posterior.predict(Xtest, train_data, covariance="dense")
+    dense_dist = dense_posterior.predict(Xtest, train_data, covariance="dense")
+
+    assert isinstance(ss_dist.scale, lx.MatrixLinearOperator)
+    np.testing.assert_allclose(
+        np.asarray(ss_dist.mean), np.asarray(dense_dist.mean), atol=1e-8
+    )
+    np.testing.assert_allclose(
+        np.asarray(ss_dist.covariance_matrix),
+        np.asarray(dense_dist.covariance_matrix),
+        atol=1e-8,
+    )
+
+
+def test_state_space_posterior_predict_smoothed_dense_diagonal_matches_diagonal_mode():
+    """The diagonal of the dense covariance must equal the diagonal-mode variances."""
+    ss_posterior, _, train_data, Xtest = _build_smoothed_dense_comparison(
+        gpx.kernels.Matern32, jitter=1e-6
+    )
+
+    dense_dist = ss_posterior.predict(Xtest, train_data, covariance="dense")
+    diagonal_dist = ss_posterior.predict(Xtest, train_data, covariance="diagonal")
+
+    np.testing.assert_allclose(
+        np.diag(np.asarray(dense_dist.covariance_matrix)),
+        np.asarray(diagonal_dist.variance),
+        atol=1e-10,
+    )
+
+
+def test_state_space_posterior_predict_smoothed_dense_preserves_caller_order():
+    """Unsorted test inputs must produce a covariance matrix permuted the same
+    way as sorted test inputs — this exercises the grid-order/caller-order
+    bookkeeping unique to the dense path."""
+    ss_posterior, _, train_data, _ = _build_smoothed_dense_comparison(
+        gpx.kernels.Matern32, jitter=1e-6
+    )
+    Xtest_caller_order = jnp.array([5.5, 0.5, 9.5, 2.5, 7.5]).reshape(-1, 1)
+    sort_perm = np.argsort(np.asarray(Xtest_caller_order).squeeze())
+    Xtest_sorted = Xtest_caller_order[sort_perm]
+
+    dist_caller = ss_posterior.predict(
+        Xtest_caller_order, train_data, covariance="dense"
+    )
+    dist_sorted = ss_posterior.predict(Xtest_sorted, train_data, covariance="dense")
+
+    cov_caller = np.asarray(dist_caller.covariance_matrix)
+    cov_sorted = np.asarray(dist_sorted.covariance_matrix)
+    np.testing.assert_allclose(
+        cov_caller,
+        cov_sorted[np.ix_(np.argsort(sort_perm), np.argsort(sort_perm))],
+        atol=1e-9,
+    )
+
+
+def test_state_space_posterior_predict_smoothed_dense_single_test_point():
+    """M=1 is the degenerate case with no cross-covariance segments at all."""
+    ss_posterior, dense_posterior, train_data, _ = _build_smoothed_dense_comparison(
+        gpx.kernels.Matern12, jitter=1e-6
+    )
+    Xtest = jnp.array([[4.2]])
+
+    ss_dist = ss_posterior.predict(Xtest, train_data, covariance="dense")
+    dense_dist = dense_posterior.predict(Xtest, train_data, covariance="dense")
+
+    assert ss_dist.covariance_matrix.shape == (1, 1)
+    np.testing.assert_allclose(
+        np.asarray(ss_dist.covariance_matrix),
+        np.asarray(dense_dist.covariance_matrix),
+        atol=1e-8,
+    )
+
+
+def test_state_space_posterior_predict_smoothed_dense_joint_sampling_smoke():
+    """Joint samples must be finite and their empirical covariance must be
+    consistent with the analytic dense covariance."""
+    ss_posterior, _, train_data, Xtest = _build_smoothed_dense_comparison(
+        gpx.kernels.Matern32, jitter=1e-6, n_test=6
+    )
+    dist = ss_posterior.predict(Xtest, train_data, covariance="dense")
+
+    samples = dist.sample(jr.key(0), sample_shape=(4000,))
+    assert samples.shape == (4000, 6)
+    assert bool(jnp.all(jnp.isfinite(samples)))
+
+    empirical_cov = np.cov(np.asarray(samples).T)
+    analytic_cov = np.asarray(dist.covariance_matrix)
+    # Monte-Carlo tolerance for 4000 draws, not the ~1e-8 exact-match bar above.
+    np.testing.assert_allclose(empirical_cov, analytic_cov, atol=0.05)
+
+
+def test_state_space_posterior_predict_smoothed_dense_is_jittable():
+    lengthscale, variance, obs_stddev = 1.2, 0.9, 0.15
+    X_train, y_train = _build_matern12_dataset(
+        n=20, lengthscale=lengthscale, variance=variance, obs_stddev=obs_stddev
+    )
+    train_data = gpx.Dataset(X=X_train.reshape(-1, 1), y=y_train.reshape(-1, 1))
+    Xtest = jnp.linspace(0.0, 10.0, 5).reshape(-1, 1)
+
+    def build_and_predict(kernel_lengthscale):
+        posterior = StateSpacePrior(
+            mean_function=gpx.mean_functions.Zero(),
+            kernel=gpx.kernels.Matern32(
+                lengthscale=kernel_lengthscale, variance=variance
+            ),
+            jitter=1e-6,
+        ) * gpx.likelihoods.Gaussian(obs_stddev=obs_stddev)
+        dist = posterior.predict(Xtest, train_data, covariance="dense")
+        return dist.mean, dist.covariance_matrix
+
+    jitted_mean, jitted_cov = jax.jit(build_and_predict)(jnp.asarray(lengthscale))
+    eager_mean, eager_cov = build_and_predict(jnp.asarray(lengthscale))
+
+    np.testing.assert_allclose(
+        np.asarray(jitted_mean), np.asarray(eager_mean), atol=1e-10
+    )
+    np.testing.assert_allclose(
+        np.asarray(jitted_cov), np.asarray(eager_cov), atol=1e-10
+    )
+
+
+def test_state_space_posterior_predict_smoothed_dense_is_differentiable():
+    lengthscale, variance, obs_stddev = 1.2, 0.9, 0.15
+    X_train, y_train = _build_matern12_dataset(
+        n=20, lengthscale=lengthscale, variance=variance, obs_stddev=obs_stddev
+    )
+    train_data = gpx.Dataset(X=X_train.reshape(-1, 1), y=y_train.reshape(-1, 1))
+    Xtest = jnp.linspace(0.0, 10.0, 5).reshape(-1, 1)
+
+    def loss(kernel_lengthscale):
+        posterior = StateSpacePrior(
+            mean_function=gpx.mean_functions.Zero(),
+            kernel=gpx.kernels.Matern32(
+                lengthscale=kernel_lengthscale, variance=variance
+            ),
+            jitter=1e-6,
+        ) * gpx.likelihoods.Gaussian(obs_stddev=obs_stddev)
+        dist = posterior.predict(Xtest, train_data, covariance="dense")
+        return jnp.sum(dist.covariance_matrix) + jnp.sum(dist.mean**2)
+
+    gradient = jax.grad(loss)(jnp.asarray(lengthscale))
+    assert jnp.isfinite(gradient)
+    assert gradient != 0.0
+
+
+def test_rts_smoother_return_gains_shape_is_linear_in_grid_size():
+    """The gains that feed the cross-covariance recursion are one array of
+    shape ``(grid_size - 1, state_dim, state_dim)`` — never an
+    ``(N, N)``-shaped object — which is what keeps the dense predictive
+    linear in the training set size rather than quadratic."""
+    from gpjax.state_space.inference import _sqrt_filter_forward
+    from gpjax.state_space.sde import Matern32SDE
+
+    lengthscale, variance, obs_stddev = 1.0, 1.0, 0.2
+    n = 40
+    X, y = _build_matern12_dataset(
+        n=n, lengthscale=lengthscale, variance=variance, obs_stddev=obs_stddev
+    )
+    sde = Matern32SDE(lengthscale=lengthscale, variance=variance)
+    time_steps = jnp.concatenate([jnp.array([0.0]), jnp.diff(X)])
+    is_observed = jnp.ones(n, dtype=bool)
+    forward_outputs, _ = _sqrt_filter_forward(
+        sde, y, time_steps, is_observed, jnp.asarray(obs_stddev)
+    )
+    _, _, smoother_gains = rts_smoother(
+        sde, forward_outputs, time_steps, return_gains=True
+    )
+    assert smoother_gains.shape == (n - 1, sde.state_dim, sde.state_dim)
+
+
+def test_state_space_posterior_predict_smoothed_dense_scales_to_larger_training_set():
+    """Soft robustness check: the dense predictive stays finite and correct
+    as the training set grows well past the small sizes used elsewhere in
+    this file, with a small, fixed number of test points."""
+    lengthscale, variance, obs_stddev = 1.0, 1.0, 0.2
+    n_train = 300
+    X_train, y_train = _build_matern12_dataset(
+        n=n_train, lengthscale=lengthscale, variance=variance, obs_stddev=obs_stddev
+    )
+    train_data = gpx.Dataset(X=X_train.reshape(-1, 1), y=y_train.reshape(-1, 1))
+    Xtest = jnp.array([[1.5], [4.5], [8.5]])
+
+    posterior = StateSpacePrior(
+        mean_function=gpx.mean_functions.Zero(),
+        kernel=gpx.kernels.Matern32(lengthscale=lengthscale, variance=variance),
+        jitter=1e-6,
+    ) * gpx.likelihoods.Gaussian(obs_stddev=obs_stddev)
+    dist = posterior.predict(Xtest, train_data, covariance="dense")
+
+    cov = np.asarray(dist.covariance_matrix)
+    assert cov.shape == (3, 3)
+    assert np.all(np.isfinite(cov))
+    np.testing.assert_allclose(cov, cov.T, atol=1e-9)
+    assert np.all(np.linalg.eigvalsh(cov) > 0.0)
+
+
 def test_state_space_posterior_predict_smoothed_returns_diagonal_distribution():
     """Returned scale must be a DiagonalLinearOperator (not dense)."""
     n = 10
diff --git a/tests/test_state_space/test_smoother.py b/tests/test_state_space/test_smoother.py
index 30a66d9bf..716953df8 100644
--- a/tests/test_state_space/test_smoother.py
+++ b/tests/test_state_space/test_smoother.py
@@ -113,6 +113,7 @@ def _numpy_rts_reference(sde, y, time_steps, obs_stddev_squared):
 
     means_smoothed = [None] * n
     covs_smoothed = [None] * n
+    gains = [None] * (n - 1)
     means_smoothed[-1] = means_filtered[-1]
     covs_smoothed[-1] = covs_filtered[-1]
     for i in range(n - 2, -1, -1):
@@ -123,6 +124,7 @@ def _numpy_rts_reference(sde, y, time_steps, obs_stddev_squared):
         smoother_gain = (
             cov_filtered @ transition_matrix_next.T @ np.linalg.inv(cov_predicted_next)
         )
+        gains[i] = smoother_gain
         means_smoothed[i] = means_filtered[i] + smoother_gain @ (
             means_smoothed[i + 1] - means_predicted[i + 1]
         )
@@ -132,7 +134,7 @@ def _numpy_rts_reference(sde, y, time_steps, obs_stddev_squared):
             @ (covs_smoothed[i + 1] - cov_predicted_next)
             @ smoother_gain.T
         )
-    return np.array(means_smoothed), np.array(covs_smoothed)
+    return np.array(means_smoothed), np.array(covs_smoothed), np.array(gains)
 
 
 def test_rts_smoother_matches_numpy_reference_to_machine_precision():
@@ -154,7 +156,7 @@ def test_rts_smoother_matches_numpy_reference_to_machine_precision():
     )
     smoothed_means, smoothed_Ls = rts_smoother(sde, forward_outputs, time_steps)
 
-    means_reference, covs_reference = _numpy_rts_reference(
+    means_reference, covs_reference, _gains_reference = _numpy_rts_reference(
         sde, y, time_steps, float(obs_stddev**2)
     )
     smoothed_covs = jax.vmap(lambda L: L @ L.T)(smoothed_Ls)
@@ -167,6 +169,49 @@ def test_rts_smoother_matches_numpy_reference_to_machine_precision():
     )
 
 
+def test_rts_smoother_return_gains_matches_numpy_reference():
+    """``return_gains=True`` exposes exactly the gains used internally; check
+    them against the same NumPy oracle used for the smoothed means/covariances,
+    and confirm the default call is unaffected (2-tuple, unchanged values)."""
+    lengthscale, variance, obs_stddev = 1.5, 0.8, 0.2
+    n = 15
+    X, y = _build_matern12_dataset(
+        n=n, lengthscale=lengthscale, variance=variance, obs_stddev=obs_stddev
+    )
+    sigma_eff = jnp.asarray(obs_stddev)
+    sde = Matern12SDE(lengthscale=lengthscale, variance=variance)
+    time_steps = jnp.concatenate([jnp.array([0.0]), jnp.diff(X)])
+    is_observed = jnp.ones(n, dtype=bool)
+
+    forward_outputs, _ = _sqrt_filter_forward(
+        sde, y, time_steps, is_observed, sigma_eff
+    )
+    smoothed_means, smoothed_Ls, smoother_gains = rts_smoother(
+        sde, forward_outputs, time_steps, return_gains=True
+    )
+    smoothed_means_default, smoothed_Ls_default = rts_smoother(
+        sde, forward_outputs, time_steps
+    )
+
+    means_reference, _covs_reference, gains_reference = _numpy_rts_reference(
+        sde, y, time_steps, float(obs_stddev**2)
+    )
+
+    assert smoother_gains.shape == (n - 1, sde.state_dim, sde.state_dim)
+    np.testing.assert_allclose(
+        np.asarray(smoother_gains), gains_reference, atol=1e-12, rtol=1e-12
+    )
+    np.testing.assert_allclose(
+        np.asarray(smoothed_means), np.asarray(smoothed_means_default), atol=0.0
+    )
+    np.testing.assert_allclose(
+        np.asarray(smoothed_Ls), np.asarray(smoothed_Ls_default), atol=0.0
+    )
+    np.testing.assert_allclose(
+        np.asarray(smoothed_means), means_reference, atol=1e-12, rtol=1e-12
+    )
+
+
 def test_smoother_is_finite_under_near_noiseless_dense_sampling():
     """Robustness guard: stiff regime (tiny obs noise, dense Matern-5/2 grid)
     must stay finite.