Repository: GPflow/GPflow

Issue:
The posterior objects work poorly with Bayesian Optimisation
# Bug

## To reproduce

**Minimal, reproducible example**

1. Create a model where the data are `tf.Variable`s.
2. Create a `Posterior` from the model.
3. Wrap `posterior.predict_f` in `tf.function`.
4. Add rows to the data.
5. Call `posterior.update_cache()`.
6. Call the compiled `predict_f`.

Observe crashes and wrong predictions.

## Expected behavior

The `Posterior` object should reflect the state of the data, even when it changes shape.


Discussion:

Resolving patch:
diff --git a/doc/source/notebooks/advanced/variational_fourier_features.pct.py b/doc/source/notebooks/advanced/variational_fourier_features.pct.py
--- a/doc/source/notebooks/advanced/variational_fourier_features.pct.py
+++ b/doc/source/notebooks/advanced/variational_fourier_features.pct.py
@@ -331,7 +331,7 @@ def _precompute(self):
 
         Qinv = Kuu.inverse().to_dense() - KuuInv_covu_KuuInv
 
-        return alpha, Qinv
+        return gpflow.posteriors.PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)
 
     def _conditional_with_precompute(self, cache, Xnew, full_cov, full_output_cov):
         alpha, Qinv = cache
diff --git a/gpflow/posteriors.py b/gpflow/posteriors.py
--- a/gpflow/posteriors.py
+++ b/gpflow/posteriors.py
@@ -14,6 +14,7 @@
 
 import enum
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import Optional, Tuple, Type, Union
 
 import tensorflow as tf
@@ -93,6 +94,57 @@ class PrecomputeCacheType(enum.Enum):
     NOCACHE = "nocache"
 
 
+@dataclass
+class PrecomputedValue:
+    value: tf.Tensor
+    """
+    The precomputed value itself.
+    """
+
+    axis_dynamic: Tuple[bool, ...]
+    """
+    A tuple with one element per dimension of `value`. That element is `True` if that dimension
+    of `value` might change size.
+    """
+
+    def __post_init__(self) -> None:
+        tf.debugging.assert_rank(
+            self.value,
+            len(self.axis_dynamic),
+            "axis_dynamic must have one element per dimension of value.",
+        )
+
+    @staticmethod
+    def wrap_alpha_Qinv(alpha: TensorType, Qinv: TensorType) -> Tuple["PrecomputedValue", ...]:
+        """
+        Wraps `alpha` and `Qinv` in `PrecomputedValue`s.
+        """
+        one_dynamic = False
+        L_dynamic = False
+        M_dynamic = False  # TODO(jesper): Support variable number of inducing points?
+
+        alpha_rank = tf.rank(alpha)
+        if alpha_rank == 2:
+            alpha_dynamic: Tuple[bool, ...] = (M_dynamic, L_dynamic)
+        elif alpha_rank == 3:
+            alpha_dynamic = (L_dynamic, M_dynamic, one_dynamic)
+        else:
+            raise AssertionError(f"Unknown rank of alpha {alpha_rank}.")
+
+        Qinv_rank = tf.rank(Qinv)
+        if Qinv_rank == 2:
+            Qinv_dynamic: Tuple[bool, ...] = (M_dynamic, M_dynamic)
+        elif Qinv_rank == 3:
+            Qinv_dynamic = (L_dynamic, M_dynamic, M_dynamic)
+        else:
+            raise AssertionError(f"Unknown rank of Qinv {Qinv_rank}.")
+
+        return (
+            PrecomputedValue(alpha, alpha_dynamic),
+            PrecomputedValue(Qinv, Qinv_dynamic),
+        )
+
+
 def _validate_precompute_cache_type(
     value: Union[None, PrecomputeCacheType, str]
 ) -> PrecomputeCacheType:
@@ -104,7 +156,8 @@ def _validate_precompute_cache_type(
         return PrecomputeCacheType(value.lower())
     else:
         raise ValueError(
-            f"{value} is not a valid PrecomputeCacheType. Valid options: 'tensor', 'variable', 'nocache' (or None)."
+            f"{value} is not a valid PrecomputeCacheType."
+            " Valid options: 'tensor', 'variable', 'nocache' (or None)."
         )
 
 
@@ -139,7 +192,7 @@ def _add_mean_function(self, Xnew: TensorType, mean: TensorType) -> tf.Tensor:
             return mean + self.mean_function(Xnew)
 
     @abstractmethod
-    def _precompute(self) -> Tuple[tf.Tensor, ...]:
+    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
         """
         Precompute a cache.
 
@@ -206,7 +259,8 @@ def update_cache(self, precompute_cache: Optional[PrecomputeCacheType] = None) -
         if precompute_cache is None:
             if self._precompute_cache is None:
                 raise ValueError(
-                    "You must pass precompute_cache explicitly (the cache had not been updated before)."
+                    "You must pass precompute_cache explicitly"
+                    " (the cache had not been updated before)."
                 )
             precompute_cache = self._precompute_cache
         else:
@@ -216,16 +270,23 @@ def update_cache(self, precompute_cache: Optional[PrecomputeCacheType] = None) -
             self.cache = None
 
         elif precompute_cache is PrecomputeCacheType.TENSOR:
-            self.cache = self._precompute()
+            self.cache = tuple(c.value for c in self._precompute())
 
         elif precompute_cache is PrecomputeCacheType.VARIABLE:
             cache = self._precompute()
+
             if self.cache is not None and all(isinstance(c, tf.Variable) for c in self.cache):
                 # re-use existing variables
-                for cache_var, new_value in zip(self.cache, cache):
-                    cache_var.assign(new_value)
+                for cache_var, c in zip(self.cache, cache):
+                    cache_var.assign(c.value)
             else:  # create variables
-                self.cache = tuple(tf.Variable(new_value, trainable=False) for new_value in cache)
+                shapes = [
+                    [None if d else s for d, s in zip(c.axis_dynamic, tf.shape(c.value))]
+                    for c in cache
+                ]
+                self.cache = tuple(
+                    tf.Variable(c.value, trainable=False, shape=s) for c, s in zip(cache, shapes)
+                )
 
 
 class GPRPosterior(AbstractPosterior):
@@ -301,12 +362,17 @@ def _conditional_with_precompute(
 
         return mean, cov
 
-    def _precompute(self) -> Tuple[tf.Tensor, ...]:
+    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
         Kmm = self.kernel(self.X_data)
         Kmm_plus_s = add_noise_cov(Kmm, self.likelihood_variance)
 
         Lm = tf.linalg.cholesky(Kmm_plus_s)
-        Kmm_plus_s_inv = tf.linalg.cholesky_solve(Lm, tf.eye(self.X_data.shape[0], dtype=Lm.dtype))
+        Kmm_plus_s_inv = tf.linalg.cholesky_solve(
+            Lm, tf.eye(tf.shape(self.X_data)[0], dtype=Lm.dtype)
+        )
+
+        M = self.X_data.shape[0]
+        M_dynamic = M is None
 
         tf.debugging.assert_shapes(
             [
@@ -314,7 +380,7 @@ def _precompute(self) -> Tuple[tf.Tensor, ...]:
                 (Kmm, ["M", "M"]),
             ]
         )
-        return (Kmm_plus_s_inv,)
+        return (PrecomputedValue(Kmm_plus_s_inv, (M_dynamic, M_dynamic)),)
 
     def _conditional_fused(
         self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
@@ -365,7 +431,7 @@ def __init__(
         if precompute_cache is not None:
             self.update_cache(precompute_cache)
 
-    def _precompute(self) -> Tuple[tf.Tensor, ...]:
+    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
         # taken directly from the deprecated SGPR implementation
         num_inducing = self.inducing_variable.num_inducing
         assert self.mean_function is not None
@@ -393,7 +459,7 @@ def _precompute(self) -> Tuple[tf.Tensor, ...]:
         alpha = LinvT @ tf.transpose(LBinv) @ c
         Qinv = LinvT @ tmp @ Linv
 
-        return alpha, Qinv
+        return PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)
 
     def _conditional_with_precompute(
         self,
@@ -512,18 +578,21 @@ def _conditional_with_precompute(
             white=self.white,
         )
 
-    def _precompute(self) -> Tuple[tf.Tensor, ...]:
+    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
         Kmm = self.kernel(self.X_data) + eye(
             tf.shape(self.X_data)[-2], value=default_jitter(), dtype=self.X_data.dtype
         )  # [..., M, M]
         Lm = tf.linalg.cholesky(Kmm)
 
-        return (Lm,)
+        M = self.X_data.shape[0]
+        M_dynamic = M is None
+
+        return (PrecomputedValue(Lm, (M_dynamic, M_dynamic)),)
 
     def _conditional_fused(
         self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
     ) -> MeanAndVariance:
-        temp_cache = self._precompute()
+        temp_cache = tuple(c.value for c in self._precompute())
         return self._conditional_with_precompute(temp_cache, Xnew, full_cov, full_output_cov)
 
 
@@ -563,7 +632,7 @@ def _set_qdist(self, q_mu: TensorType, q_sqrt: TensorType) -> tf.Tensor:
         else:
             self._q_dist = _MvNormal(q_mu, q_sqrt)
 
-    def _precompute(self) -> Tuple[tf.Tensor, ...]:
+    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
         Kuu = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [(R), M, M]
         q_mu = self._q_dist.q_mu
 
@@ -621,7 +690,7 @@ def _precompute(self) -> Tuple[tf.Tensor, ...]:
             ]
         )
 
-        return alpha, Qinv
+        return PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)
 
 
 class IndependentPosterior(BasePosterior):
@@ -635,14 +704,16 @@ def _get_Kff(self, Xnew: TensorType, full_cov: bool) -> tf.Tensor:
         # TODO: this assumes that Xnew has shape [N, D] and no leading dims
 
         if isinstance(self.kernel, (kernels.SeparateIndependent, kernels.IndependentLatent)):
-            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would return
+            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
+            # return
             # if full_cov: [P, N, N] -- this is what we want
             # else: [N, P] instead of [P, N] as we get from the explicit stack below
             Kff = tf.stack([k(Xnew, full_cov=full_cov) for k in self.kernel.kernels], axis=0)
         elif isinstance(self.kernel, kernels.MultioutputKernel):
             # effectively, SharedIndependent path
             Kff = self.kernel.kernel(Xnew, full_cov=full_cov)
-            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would return
+            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
+            # return
             # if full_cov: [P, N, N] instead of [N, N]
             # else: [N, P] instead of [N]
         else:
@@ -687,7 +758,8 @@ class IndependentPosteriorSingleOutput(IndependentPosterior):
     def _conditional_fused(
         self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
     ) -> MeanAndVariance:
-        # same as IndependentPosteriorMultiOutput, Shared~/Shared~ branch, except for following line:
+        # same as IndependentPosteriorMultiOutput, Shared~/Shared~ branch, except for following
+        # line:
         Knn = self.kernel(Xnew, full_cov=full_cov)
 
         Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
@@ -717,7 +789,8 @@ def _conditional_fused(
                 Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
             )  # [N, P],  [P, N, N] or [N, P]
         else:
-            # this is the messy thing with tf.map_fn, cleaned up by the st/clean_up_broadcasting_conditionals branch
+            # this is the messy thing with tf.map_fn, cleaned up by the
+            # st/clean_up_broadcasting_conditionals branch
 
             # Following are: [P, M, M]  -  [P, M, N]  -  [P, N](x N)
             Kmms = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [P, M, M]


Test patch:
diff --git a/tests/gpflow/posteriors/conftest.py b/tests/gpflow/posteriors/conftest.py
new file mode 100644
--- /dev/null
+++ b/tests/gpflow/posteriors/conftest.py
@@ -0,0 +1,57 @@
+#  Copyright 2022 The GPflow Contributors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from inspect import isabstract
+from typing import DefaultDict, Iterable, Set, Type
+
+import pytest
+
+import gpflow.ci_utils
+from gpflow.posteriors import AbstractPosterior
+
+
+@pytest.fixture(name="tested_posteriors", scope="package")
+def _tested_posteriors() -> DefaultDict[str, Set[Type[AbstractPosterior]]]:
+    return DefaultDict(set)
+
+
+@pytest.fixture(scope="package", autouse=True)
+def _ensure_all_posteriors_are_tested_fixture(
+    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]]
+) -> Iterable[None]:
+    """
+    This fixture ensures that all concrete posteriors have unit tests which compare the predictions
+    from the fused and precomputed code paths. When adding a new concrete posterior class to
+    GPFlow, ensure that it is also tested in this manner.
+
+    This autouse, package scoped fixture will always be executed when tests in this package are run.
+    """
+    # Code here will be executed before any of the tests in this package.
+
+    yield  # Run tests in this package.
+
+    # Code here will be executed after all of the tests in this package.
+
+    available_posteriors = list(gpflow.ci_utils.subclasses(AbstractPosterior))
+    concrete_posteriors = set([k for k in available_posteriors if not isabstract(k)])
+
+    messages = []
+    for key, key_tested_posteriors in tested_posteriors.items():
+        untested_posteriors = concrete_posteriors - key_tested_posteriors
+        if untested_posteriors:
+            messages.append(
+                f"For key '{key}' no tests have been registered for the following posteriors: {untested_posteriors}."
+            )
+
+    if messages:
+        raise AssertionError("\n".join(messages))
diff --git a/tests/gpflow/posteriors/test_bo_integration.py b/tests/gpflow/posteriors/test_bo_integration.py
new file mode 100644
--- /dev/null
+++ b/tests/gpflow/posteriors/test_bo_integration.py
@@ -0,0 +1,403 @@
+#  Copyright 2022 The GPflow Contributors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from typing import Any, Callable, DefaultDict, Dict, Iterator, Mapping, Set, Tuple, Type, TypeVar
+
+import numpy as np
+import pytest
+import tensorflow as tf
+from _pytest.fixtures import SubRequest
+
+import gpflow
+from gpflow.base import RegressionData
+from gpflow.config import default_float
+from gpflow.inducing_variables import InducingPoints, InducingVariables
+from gpflow.kernels import Kernel, Matern52
+from gpflow.likelihoods import Exponential, Likelihood
+from gpflow.models import GPR, SGPR, SVGP, VGP, GPModel, training_loss_closure
+from gpflow.models.vgp import update_vgp_data
+from gpflow.posteriors import AbstractPosterior, PrecomputeCacheType
+
+_CreateModel = Callable[[RegressionData], GPModel]
+_C = TypeVar("_C", bound=_CreateModel)
+
+_MULTI_OUTPUT = "multi_output"
+_MODEL_FACTORIES: Dict[_CreateModel, Mapping[str, Any]] = {}
+
+# This exists to make it easy to disable tf.function, for debugging.
+_COMPILE = True
+_MAXITER = 500
+_DEFAULT_ATOL = 1e-10
+_DEFAULT_RTOL = 1e-7
+
+
+@pytest.fixture(name="register_posterior_bo_integration_test")
+def _register_posterior_bo_integration_test(
+    request: SubRequest,
+    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]],
+) -> Callable[[AbstractPosterior], None]:
+    def _register_posterior(posterior: AbstractPosterior) -> None:
+        tested_posteriors[request.function.__name__].add(posterior.__class__)
+
+    return _register_posterior
+
+
+def model_factory(
+    *flags: str, atol: float = _DEFAULT_ATOL, rtol: float = _DEFAULT_RTOL
+) -> Callable[[_C], _C]:
+    """ Decorator for adding a function to the `_MODEL_FACTORIES` list. """
+
+    properties = {
+        "atol": atol,
+        "rtol": rtol,
+        **{flag: True for flag in flags},
+    }
+
+    def register(create_model: _C) -> _C:
+        _MODEL_FACTORIES[create_model] = properties
+        return create_model
+
+    return register
+
+
+def create_kernel() -> Kernel:
+    return Matern52()
+
+
+def create_likelihood() -> Likelihood:
+    return Exponential()
+
+
+def create_inducing_points(data: RegressionData) -> InducingPoints:
+    n_features = data[0].shape[1]
+    n_inducing_points = 25
+    rng = np.random.default_rng(20220208)
+    Z = tf.constant(rng.random((n_inducing_points, n_features)))
+    return InducingPoints(Z)
+
+
+def create_q(
+    inducing_variable: InducingVariables, *, row_scale: int = 1, column_scale: int = 1
+) -> Tuple[bool, tf.Tensor, tf.Tensor]:
+    n_inducing_points = inducing_variable.num_inducing
+    rng = np.random.default_rng(20220133)
+    q_diag = True
+    q_mu = tf.constant(rng.random((row_scale * n_inducing_points, column_scale)))
+    q_sqrt = tf.constant(rng.random((row_scale * n_inducing_points, column_scale))) ** 2
+    return q_diag, q_mu, q_sqrt
+
+
+@model_factory(rtol=1e-3)
+def create_gpr(data: RegressionData) -> GPR:
+    return GPR(data=data, kernel=create_kernel())
+
+
+@model_factory(rtol=1e-4)
+def create_sgpr(data: RegressionData) -> SGPR:
+    return SGPR(data=data, kernel=create_kernel(), inducing_variable=create_inducing_points(data))
+
+
+@model_factory(rtol=1e-3)
+def create_vgp(data: RegressionData) -> VGP:
+    return VGP(data=data, kernel=create_kernel(), likelihood=create_likelihood())
+
+
+@model_factory()
+def create_svgp__independent_single_output(data: RegressionData) -> SVGP:
+    inducing_variable = create_inducing_points(data)
+    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
+    return SVGP(
+        kernel=create_kernel(),
+        likelihood=create_likelihood(),
+        inducing_variable=inducing_variable,
+        q_diag=q_diag,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+
+
+@model_factory(_MULTI_OUTPUT)
+def create_svgp__fully_correlated_multi_output(data: RegressionData) -> SVGP:
+    n_outputs = data[1].shape[1]
+    kernel = gpflow.kernels.SharedIndependent(create_kernel(), output_dim=n_outputs)
+    inducing_variable = create_inducing_points(data)
+    q_diag, q_mu, q_sqrt = create_q(inducing_variable, row_scale=n_outputs)
+    return SVGP(
+        kernel=kernel,
+        likelihood=create_likelihood(),
+        inducing_variable=inducing_variable,
+        q_diag=q_diag,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+
+
+@model_factory(_MULTI_OUTPUT)
+def create_svgp__independent_multi_output(data: RegressionData) -> SVGP:
+    n_outputs = data[1].shape[1]
+    kernel = gpflow.kernels.SharedIndependent(create_kernel(), output_dim=n_outputs)
+    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
+        create_inducing_points(data)
+    )
+    q_diag, q_mu, q_sqrt = create_q(inducing_variable, column_scale=n_outputs)
+    return SVGP(
+        kernel=kernel,
+        likelihood=create_likelihood(),
+        inducing_variable=inducing_variable,
+        q_diag=q_diag,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+
+
+@model_factory(_MULTI_OUTPUT)
+def create_svgp__fallback_independent_latent_posterior(data: RegressionData) -> SVGP:
+    n_outputs = data[1].shape[1]
+    rng = np.random.default_rng(20220131)
+    kernel = gpflow.kernels.LinearCoregionalization(
+        [create_kernel()],
+        W=tf.constant(rng.standard_normal((n_outputs, 1))),
+    )
+    inducing_variable = gpflow.inducing_variables.FallbackSeparateIndependentInducingVariables(
+        [create_inducing_points(data)]
+    )
+    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
+    return SVGP(
+        kernel=kernel,
+        likelihood=create_likelihood(),
+        inducing_variable=inducing_variable,
+        q_diag=q_diag,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+
+
+@model_factory(_MULTI_OUTPUT)
+def create_svgp__linear_coregionalization(data: RegressionData) -> SVGP:
+    n_outputs = data[1].shape[1]
+    rng = np.random.default_rng(20220131)
+    kernel = gpflow.kernels.LinearCoregionalization(
+        [create_kernel()], W=tf.constant(rng.standard_normal((n_outputs, 1)))
+    )
+    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
+        create_inducing_points(data)
+    )
+    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
+    return SVGP(
+        kernel=kernel,
+        likelihood=create_likelihood(),
+        inducing_variable=inducing_variable,
+        q_diag=q_diag,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+
+
+@pytest.fixture(params=_MODEL_FACTORIES)
+def _create_model(request: SubRequest) -> _CreateModel:
+    return request.param
+
+
+@pytest.fixture
+def _multi_output(_create_model: _CreateModel) -> bool:
+    return _MULTI_OUTPUT in _MODEL_FACTORIES[_create_model]
+
+
+@pytest.fixture
+def _rtol(_create_model: _CreateModel) -> float:
+    return _MODEL_FACTORIES[_create_model]["rtol"]
+
+
+@pytest.fixture
+def _atol(_create_model: _CreateModel) -> float:
+    return _MODEL_FACTORIES[_create_model]["atol"]
+
+
+@pytest.fixture
+def _f_minimum(_multi_output: bool) -> tf.Tensor:
+    return (
+        tf.constant(
+            [
+                [0.2, 0.4],
+                [0.4, 0.6],
+                [0.6, 0.8],
+            ],
+            dtype=default_float(),
+        )
+        if _multi_output
+        else tf.constant([[0.3, 0.5]], dtype=default_float())
+    )
+
+
+@pytest.fixture
+def _f(_f_minimum: tf.Tensor) -> Callable[[tf.Tensor], tf.Tensor]:
+    def f(X: tf.Tensor) -> tf.Tensor:
+        err = X[:, None, :] - _f_minimum[None, :, :]
+        err_sq = err ** 2
+        return tf.reduce_sum(err_sq, axis=-1)
+
+    return f
+
+
+@pytest.fixture
+def _data(
+    _f: Callable[[tf.Tensor], tf.Tensor], _f_minimum: tf.Tensor
+) -> Tuple[tf.Variable, tf.Variable]:
+    n_initial_data = 10
+    n_outputs, n_features = _f_minimum.shape
+
+    rng = np.random.default_rng(20220126)
+    X = tf.Variable(
+        rng.random((n_initial_data, n_features)),
+        shape=[None, n_features],
+        dtype=default_float(),
+        trainable=False,
+    )
+    Y = tf.Variable(
+        _f(X),
+        shape=[None, n_outputs],
+        dtype=default_float(),
+        trainable=False,
+    )
+
+    return X, Y
+
+
+@pytest.fixture
+def _extend_data(
+    _data: Tuple[tf.Variable, tf.Variable], _f: Callable[[tf.Tensor], tf.Tensor]
+) -> Callable[[GPModel], Iterator[int]]:
+    n_iterations = 3
+    rng = np.random.default_rng(20220127)
+    X, Y = _data
+    n_features = X.shape[1]
+
+    def iterate(model: GPModel) -> Iterator[int]:
+        for i in range(n_iterations):
+            X_new = tf.constant(rng.random((1, n_features)))
+            Y_new = _f(X_new)
+            X_i = tf.concat([X, X_new], axis=0)
+            Y_i = tf.concat([Y, Y_new], axis=0)
+
+            if isinstance(model, VGP):
+                update_vgp_data(model, (X_i, Y_i))
+            else:
+                X.assign(X_i)
+                Y.assign(Y_i)
+            yield i
+
+    return iterate
+
+
+@pytest.fixture
+def _X_new(_data: Tuple[tf.Variable, tf.Variable]) -> tf.Tensor:
+    rng = np.random.default_rng(20220128)
+    X, _Y = _data
+    n_features = X.shape[1]
+    return tf.constant(rng.random((3, n_features)))
+
+
+@pytest.fixture
+def _optimize(_data: Tuple[tf.Variable, tf.Variable]) -> Callable[[GPModel], None]:
+    def optimize(model: GPModel) -> None:
+        gpflow.optimizers.Scipy().minimize(
+            training_loss_closure(model, _data, compile=_COMPILE),
+            variables=model.trainable_variables,
+            options=dict(maxiter=_MAXITER),
+            method="BFGS",
+            compile=_COMPILE,
+        )
+
+    return optimize
+
+
+def test_posterior_bo_integration__predict_f(
+    register_posterior_bo_integration_test: Callable[[AbstractPosterior], None],
+    _create_model: _CreateModel,
+    _data: Tuple[tf.Variable, tf.Variable],
+    _extend_data: Callable[[GPModel], Iterator[int]],
+    _X_new: tf.Tensor,
+    _rtol: float,
+    _atol: float,
+) -> None:
+    """
+    Check that data added incrementally is correctly reflected in `predict_f`.
+    """
+    _X, Y = _data
+    n_rows_new = _X_new.shape[0]
+    n_outputs = Y.shape[1]
+
+    model = _create_model(_data)
+    posterior = model.posterior(PrecomputeCacheType.VARIABLE)
+    register_posterior_bo_integration_test(posterior)
+    predict_f = posterior.predict_f
+    if _COMPILE:
+        predict_f = tf.function(predict_f)
+
+    for _ in _extend_data(model):
+        posterior.update_cache()
+        compiled_mean, compiled_var = predict_f(_X_new)
+
+        np.testing.assert_equal((n_rows_new, n_outputs), compiled_mean.shape)
+        np.testing.assert_equal((n_rows_new, n_outputs), compiled_var.shape)
+
+        eager_model = _create_model(_data)
+        eager_mean, eager_var = eager_model.predict_f(_X_new)
+
+        np.testing.assert_allclose(eager_mean, compiled_mean, rtol=_rtol, atol=_atol)
+        np.testing.assert_allclose(eager_var, compiled_var, rtol=_rtol, atol=_atol)
+
+
+def test_posterior_bo_integration__optimization(
+    register_posterior_bo_integration_test: Callable[[AbstractPosterior], None],
+    _create_model: _CreateModel,
+    _data: Tuple[tf.Variable, tf.Variable],
+    _extend_data: Callable[[GPModel], Iterator[int]],
+    _X_new: tf.Tensor,
+    _optimize: Callable[[GPModel], None],
+    _rtol: float,
+    _atol: float,
+) -> None:
+    """
+    Check that data added incrementally is considered when optimizing a model.
+    """
+    _X, Y = _data
+    n_rows_new = _X_new.shape[0]
+    n_outputs = Y.shape[1]
+
+    model = _create_model(_data)
+    posterior = model.posterior(PrecomputeCacheType.VARIABLE)
+    register_posterior_bo_integration_test(posterior)
+    predict_f = posterior.predict_f
+    if _COMPILE:
+        predict_f = tf.function(predict_f)
+
+    # Add all the data first, and then `optimize`, so that both models are optimized the same number
+    # of times and with the same data, so they converge to the same result.
+
+    for _ in _extend_data(model):
+        pass
+
+    _optimize(model)
+    posterior.update_cache()
+    compiled_mean, compiled_var = predict_f(_X_new)
+
+    np.testing.assert_equal((n_rows_new, n_outputs), compiled_mean.shape)
+    np.testing.assert_equal((n_rows_new, n_outputs), compiled_var.shape)
+
+    eager_model = _create_model(_data)
+    _optimize(eager_model)
+    eager_mean, eager_var = eager_model.predict_f(_X_new)
+
+    np.testing.assert_allclose(eager_mean, compiled_mean, rtol=_rtol, atol=_atol)
+    np.testing.assert_allclose(eager_var, compiled_var, rtol=_rtol, atol=_atol)
diff --git a/tests/gpflow/posteriors/test_posteriors.py b/tests/gpflow/posteriors/test_posteriors.py
--- a/tests/gpflow/posteriors/test_posteriors.py
+++ b/tests/gpflow/posteriors/test_posteriors.py
@@ -13,7 +13,7 @@
 #  limitations under the License.
 import warnings
 from inspect import isabstract
-from typing import Any, Callable, Iterable, Optional, Set, Type
+from typing import Any, Callable, DefaultDict, Iterable, Optional, Set, Type
 
 import numpy as np
 import pytest
@@ -46,6 +46,25 @@
 NUM_INDUCING_POINTS = 3
 
 
+# `PosteriorType` really should be something like `Type[AbstractPosterior]`, except mypy doesn't
+# allow passing abstract classes to functions. See: https://github.com/python/mypy/issues/4717
+PosteriorType = Type[Any]
+RegisterPosterior = Callable[[AbstractPosterior, PosteriorType], None]
+
+
+@pytest.fixture(name="register_posterior_test")
+def _register_posterior_test_fixture(
+    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]]
+) -> RegisterPosterior:
+    def _verify_and_register_posterior_test(
+        posterior: AbstractPosterior, expected_posterior_class: Type[AbstractPosterior]
+    ) -> None:
+        assert isinstance(posterior, expected_posterior_class)
+        tested_posteriors["test_posteriors.py"].add(expected_posterior_class)
+
+    return _verify_and_register_posterior_test
+
+
 QSqrtFactory = Callable[[int, int], Optional[TensorType]]
 
 
@@ -94,56 +113,6 @@ def _output_dims_fixture(request: SubRequest) -> int:
     return request.param
 
 
-TESTED_POSTERIORS: Set[Type[AbstractPosterior]] = set()
-
-
-@pytest.fixture(scope="module", autouse=True)
-def _ensure_all_posteriors_are_tested_fixture() -> Iterable[None]:
-    """
-    This fixture ensures that all concrete posteriors have unit tests which compare the predictions
-    from the fused and precomputed code paths. When adding a new concrete posterior class to
-    GPFlow, ensure that it is also tested in this manner.
-
-    This autouse, module scoped fixture will always be executed when tests in this module are run.
-    """
-    # Code here will be executed before any of the tests in this module.
-
-    yield  # Run tests in this module.
-
-    # Code here will be executed after all of the tests in this module.
-
-    available_posteriors = list(gpflow.ci_utils.subclasses(AbstractPosterior))
-    concrete_posteriors = set([k for k in available_posteriors if not isabstract(k)])
-
-    untested_posteriors = concrete_posteriors - TESTED_POSTERIORS
-
-    if untested_posteriors:
-        message = (
-            f"No tests have been registered for the following posteriors: {untested_posteriors}."
-        )
-        if gpflow.ci_utils.is_continuous_integration():
-            raise AssertionError(message)
-        else:
-            warnings.warn(message)
-
-
-# `PosteriorType` really should be something like `Type[AbstractPosterior]`, except mypy doesn't
-# allow passing abstract classes to functions. See: https://github.com/python/mypy/issues/4717
-PosteriorType = Type[Any]
-RegisterPosterior = Callable[[AbstractPosterior, PosteriorType], None]
-
-
-@pytest.fixture(name="register_posterior_test")
-def _register_posterior_test_fixture() -> RegisterPosterior:
-    def _verify_and_register_posterior_test(
-        posterior: AbstractPosterior, expected_posterior_class: Type[AbstractPosterior]
-    ) -> None:
-        assert isinstance(posterior, expected_posterior_class)
-        TESTED_POSTERIORS.add(expected_posterior_class)
-
-    return _verify_and_register_posterior_test
-
-
 ConditionalClosure = Callable[..., tf.Tensor]