Repository: GPflow/GPflow

Issue:
Modelling linear noise in GPR
# Feature request

Enable modelling heteroskedastic noise in a GPR setting

## Motivation

Significant gains in performance were found across regression tasks when introducing a linear noise kernel (see Fig 1 of https://arxiv.org/abs/2106.08185). 

## Proposal

**Describe the solution you would like**

The option for a GPR model to include an input-dependent noise term. 

**What alternatives have you considered?**
A direct implementation of the linear noise kernel.

**Are you willing to open a pull request?** (We really appreciate contributions!)
Yes



Discussion:

Resolving patch:
diff --git a/doc/sphinx/notebooks/advanced/varying_noise.pct.py b/doc/sphinx/notebooks/advanced/varying_noise.pct.py
--- a/doc/sphinx/notebooks/advanced/varying_noise.pct.py
+++ b/doc/sphinx/notebooks/advanced/varying_noise.pct.py
@@ -24,284 +24,367 @@
 # y_i | f, x_i &\sim \mathcal N\big(y_i; f(x_i), \sigma^2_i\big)
 # \end{align}
 #
-# We'll demonstrate two methods. In the first demonstration, we'll assume that the noise variance is known for every data point. We'll incorporate the known noise variances $\sigma^2_i$ into the data matrix $\mathbf Y$, make a likelihood that can deal with this structure, and implement inference using variational GPs with natural gradients.
-#
-# In the second demonstration, we'll assume that the noise variance is not known, but we'd like to estimate it for different groups of data. We'll show how to construct an appropriate likelihood for this task and set up inference similarly to the first demonstration, with optimization over the noise variances.
-#
+# We'll demonstrate three methods for specifying the data-point specific noise:
+# * First we'll show how to fit the noise variance to a simple function.
+# * In the second demonstration we'll assume that the data comes from two different groups and show how to learn separate noise variance for the groups.
+# * Third we'll assume you have multiple samples at each location $x$ and show how to use the empirical variance.
+
+# %%
+import os
+import warnings
+
+warnings.simplefilter("ignore")
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
 # %%
+from typing import Optional, Tuple
+
+import matplotlib.pyplot as plt
 import numpy as np
 import tensorflow as tf
-import gpflow
+
+import gpflow as gf
 from gpflow.ci_utils import reduce_in_tests
-from gpflow.optimizers import NaturalGradient
-from gpflow import set_trainable
-import matplotlib.pyplot as plt
+from gpflow.experimental.check_shapes import inherit_check_shapes
+from gpflow.utilities import print_summary
 
 # %matplotlib inline
 
+gf.config.set_default_summary_fmt("notebook")
+
+optimizer_config = dict(maxiter=reduce_in_tests(1000))
+n_data = reduce_in_tests(300)
+X_plot = np.linspace(0.0, 1.0, reduce_in_tests(101))[:, None]
+
+
+# %% [markdown]
+# To help us later we'll first define a small function for plotting data with predictions:
+
+# %%
+def plot_distribution(
+    X: gf.base.AnyNDArray,
+    Y: gf.base.AnyNDArray,
+    X_plot: Optional[gf.base.AnyNDArray] = None,
+    mean_plot: Optional[gf.base.AnyNDArray] = None,
+    var_plot: Optional[gf.base.AnyNDArray] = None,
+    X_err: Optional[gf.base.AnyNDArray] = None,
+    mean_err: Optional[gf.base.AnyNDArray] = None,
+    var_err: Optional[gf.base.AnyNDArray] = None,
+) -> None:
+    plt.figure(figsize=(15, 5))
+
+    X = X.squeeze(axis=-1)
+    Y = Y.squeeze(axis=-1)
+    plt.scatter(X, Y, color="gray", alpha=0.8)
+
+    def get_confidence_bounds(
+        mean: gf.base.AnyNDArray, var: gf.base.AnyNDArray
+    ) -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
+        std = np.sqrt(var)
+        return mean - 1.96 * std, mean + 1.96 * std
+
+    if X_plot is not None:
+        assert mean_plot is not None
+        assert var_plot is not None
+        X_plot = X_plot.squeeze(axis=-1)
+        mean_plot = mean_plot.squeeze(axis=-1)
+        var_plot = var_plot.squeeze(axis=-1)
+        lower_plot, upper_plot = get_confidence_bounds(mean_plot, var_plot)
+        plt.fill_between(X_plot, lower_plot, upper_plot, color="silver", alpha=0.25)
+        plt.plot(X_plot, lower_plot, color="silver")
+        plt.plot(X_plot, upper_plot, color="silver")
+        plt.plot(X_plot, mean_plot, color="black")
+
+    if X_err is not None:
+        assert mean_err is not None
+        assert var_err is not None
+        lower_err, upper_err = get_confidence_bounds(mean_err, var_err)
+        plt.vlines(X_err, lower_err, upper_err, color="black")
+
+    plt.xlim([0, 1])
+    plt.show()
+
+
 # %% [markdown]
 # ## Demo 1: known noise variances
-# ### Generate synthetic data
+# ### Generate data
 # We create a utility function to generate synthetic data, including noise that varies amongst the data:
 
 # %%
-np.random.seed(1)  # for reproducibility
+def generate_data() -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
+    rng = np.random.default_rng(42)  # for reproducibility
+    X = rng.uniform(0.0, 1.0, (n_data, 1))
+    signal = (X - 0.5) ** 2 + 0.05
+    noise_scale = 0.5 * signal
+    noise = noise_scale * rng.standard_normal((n_data, 1))
+    Y = signal + noise
+    return X, Y
 
 
-def generate_data(N=80):
-    X = np.random.rand(N)[:, None] * 10 - 5  # Inputs, shape N x 1
-    F = 2.5 * np.sin(6 * X) + np.cos(3 * X)  # Mean function values
-    NoiseVar = 2 * np.exp(-((X - 2) ** 2) / 4) + 0.3  # Noise variances
-    Y = F + np.random.randn(N, 1) * np.sqrt(NoiseVar)  # Noisy data
-    return X, Y, NoiseVar
+X, Y = generate_data()
 
+# %% [markdown]
+# The data alone looks like:
 
-X, Y, NoiseVar = generate_data()
+# %%
+plot_distribution(X, Y)
 
 # %% [markdown]
-# Here's a plot of the data, with error bars representing two standard deviations:
+# ### Try a naive fit
+# If we try to fit a naive GPR model to this data, which assumes a single shared noise variance value for all data points - as specified by the likelihood variance parameter - we get:
 
 # %%
-fig, ax = plt.subplots(1, 1, figsize=(12, 6))
-_ = ax.errorbar(
-    X.squeeze(),
-    Y.squeeze(),
-    yerr=2 * (np.sqrt(NoiseVar)).squeeze(),
-    marker="x",
-    lw=0,
-    elinewidth=1.0,
-    color="C1",
+model = gf.models.GPR(
+    data=(X, Y),
+    kernel=gf.kernels.SquaredExponential(),
 )
-
-# %% [markdown]
-# ### Make a Y matrix that includes the variances
-# We need to tell the GP model what the variance is for each data point. To do this, we'll concatenate the observations with the variances into a single data matrix:
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
+)
+print_summary(model)
 
 # %%
-Y_data = np.hstack([Y, NoiseVar])
+mean, var = model.predict_y(X_plot)
+plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())
 
+# %% [markdown]
+# Notice how this naive model underestimates the noise near the center of the figure, but overestimates at for small and large values of x?
 
 # %% [markdown]
-# ### Make a new likelihood
+# ### Fit a polynomial to the noise scale
 #
-# To cope with this data structure, we'll build a new likelihood. Note how the code extracts the observations `Y` and the variances `NoiseVar` from the data. For more information on creating new likelihoods, see [Likelihood design](../tailor/likelihood_design.ipynb). Here, we're implementing the `log_prob` function (which computes the log-probability of the data given the latent function) and `variational_expectations`, which computes the expected log-probability under a Gaussian distribution on the function, and is needed in the evaluation of the evidence lower bound (ELBO). Check out the docstring for the `Likelihood` object for more information on what these functions do.
+# To fit a model with varying noise, instead of using the default single shared noise variance, we can create a Gaussian `Likelihood` with an input dependent (polynomial) `Function` for the scale of the noise,
+# then pass that likelihood to our model.
 
 # %%
-class HeteroskedasticGaussian(gpflow.likelihoods.Likelihood):
-    def __init__(self, **kwargs):
-        # this likelihood expects a single latent function F, and two columns in the data matrix Y:
-        super().__init__(input_dim=1, latent_dim=1, observation_dim=2, **kwargs)
-
-    def _log_prob(self, X, F, Y):
-        # log_prob is used by the quadrature fallback of variational_expectations and predict_log_density.
-        # Because variational_expectations is implemented analytically below, this is not actually needed,
-        # but is included for pedagogical purposes.
-        # Note that currently relying on the quadrature would fail due to https://github.com/GPflow/GPflow/issues/966
-        Y, NoiseVar = Y[:, 0], Y[:, 1]
-        return gpflow.logdensities.gaussian(Y, F, NoiseVar)
-
-    def _variational_expectations(self, X, Fmu, Fvar, Y):
-        Y, NoiseVar = Y[:, 0], Y[:, 1]
-        Fmu, Fvar = Fmu[:, 0], Fvar[:, 0]
-        return (
-            -0.5 * np.log(2 * np.pi)
-            - 0.5 * tf.math.log(NoiseVar)
-            - 0.5 * (tf.math.square(Y - Fmu) + Fvar) / NoiseVar
-        )
-
-    # The following two methods are abstract in the base class.
-    # They need to be implemented even if not used.
-
-    def _predict_log_density(self, X, Fmu, Fvar, Y):
-        raise NotImplementedError
+model = gf.models.GPR(
+    data=(X, Y),
+    kernel=gf.kernels.SquaredExponential(),
+    likelihood=gf.likelihoods.Gaussian(scale=gf.functions.Polynomial(degree=2)),
+)
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
+)
+print_summary(model)
 
-    def _predict_mean_and_var(self, X, Fmu, Fvar):
-        raise NotImplementedError
+# %%
+mean, var = model.predict_y(X_plot)
+plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())
 
 
 # %% [markdown]
-# ### Put it together with Variational Gaussian Process (VGP)
-# Here we'll build a variational GP model with the previous likelihood on the dataset that we generated. We'll use the natural gradient optimizer (see [Natural gradients](natural_gradients.ipynb) for more information).
+# ## Demo 2: grouped noise variances
 #
-# The variational GP object is capable of variational inference with any GPflow-derived likelihood. Usually, the inference is an inexact (but pretty good) approximation, but in the special case considered here, where the noise is Gaussian, it will achieve exact inference. Optimizing over the variational parameters is easy using the natural gradients method, which provably converges in a single step.
+# In this demo, we won't assume that the noise variances is a function of $x$, but we will assume that they're known to be in two groups. This example represents a case where we might know that data has been collected by two instruments with different fidelity, but we do not know what those fidelities are.
 #
-# Note: We mark the variational parameters as not trainable so that they are not included in the `model.trainable_variables` when we optimize using the Adam optimizer. We train the variational parameters separately using the natural gradient method.
-
-# %%
-# model construction
-likelihood = HeteroskedasticGaussian()
-kernel = gpflow.kernels.Matern52(lengthscales=0.5)
-model = gpflow.models.VGP((X, Y_data), kernel=kernel, likelihood=likelihood, num_latent_gps=1)
-
+# Of course it would be straightforward to add more groups. We'll stick with two for simplicity.
 
 # %% [markdown]
-# Notice that we specify num_latent_gps=1, as the VGP model would normally infer this from the shape of Y_data, but we handle the second column manually in the HeteroskedasticGaussian likelihood.
+# ### Generate data
 
 # %%
-natgrad = NaturalGradient(gamma=1.0)
-adam = tf.optimizers.Adam()
+def get_group(X: gf.base.AnyNDArray) -> gf.base.AnyNDArray:
+    return np.where((((0.2 < X) & (X < 0.4)) | ((0.7 < X) & (X < 0.8))), 0, 1)
 
-set_trainable(model.q_mu, False)
-set_trainable(model.q_sqrt, False)
 
-for _ in range(reduce_in_tests(1000)):
-    natgrad.minimize(model.training_loss, [(model.q_mu, model.q_sqrt)])
-    adam.minimize(model.training_loss, model.trainable_variables)
+def generate_grouped_data() -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
+    rng = np.random.default_rng(42)  # for reproducibility
+    X = rng.uniform(0.0, 1.0, (n_data, 1))
+    signal = np.sin(6 * X)
+    noise_scale = 0.1 + 0.4 * get_group(X)
+    noise = noise_scale * rng.standard_normal((n_data, 1))
+    Y = signal + noise
+    return X, Y
 
-# %%
-# let's do some plotting!
-xx = np.linspace(-5, 5, 200)[:, None]
-
-mu, var = model.predict_f(xx)
-
-plt.figure(figsize=(12, 6))
-plt.plot(xx, mu, "C0")
-plt.plot(xx, mu + 2 * np.sqrt(var), "C0", lw=0.5)
-plt.plot(xx, mu - 2 * np.sqrt(var), "C0", lw=0.5)
-
-plt.errorbar(
-    X.squeeze(),
-    Y.squeeze(),
-    yerr=2 * (np.sqrt(NoiseVar)).squeeze(),
-    marker="x",
-    lw=0,
-    elinewidth=1.0,
-    color="C1",
-)
-_ = plt.xlim(-5, 5)
 
-# %% [markdown]
-# ### Questions for the reader
-# 1) What is the difference in meaning between the orange vertical bars and the blue regions in the prediction?
-#
-# 2) Why did we not implement `_conditional_mean` and `_conditional_variance` in the HeteroskedasticGaussian likelihood? What could be done here?
-#
-# 2) What are some better kernel settings for this dataset? How could they be estimated?
+X, Y = generate_grouped_data()
 
 # %% [markdown]
-# ## Demo 2: grouped noise variances
-#
-# In this demo, we won't assume that the noise variances are known, but we will assume that they're known in two groups. This example represents a case where we might know that an instrument has varying fidelity for different regions, but we do not know what those fidelities are.
-#
-# Of course it would be straightforward to add more groups, or even one group per data point. We'll stick with two for simplicity.
+# And here's a plot of the raw data:
 
 # %%
-np.random.seed(1)  # for reproducibility and to make it independent from demo 1
-
+plot_distribution(X, Y)
 
 # %% [markdown]
-# ### Generate data
+# ### Fit a naive model
+#
+# Again we'll start by fitting a naive GPR, to demonstrate the problem.
 
 # %%
-def generate_data(N=100):
-    X = np.random.rand(N)[:, None] * 10 - 5  # Inputs, shape N x 1
-    F = 2.5 * np.sin(6 * X) + np.cos(3 * X)  # Mean function values
-    groups = np.where(X > 0, 0, 1)
-    NoiseVar = np.array([0.02, 0.5])[groups]  # Different variances for the two groups
-    Y = F + np.random.randn(N, 1) * np.sqrt(NoiseVar)  # Noisy data
-    return X, Y, groups
-
-
-X, Y, groups = generate_data()
+model = gf.models.GPR(
+    data=(X, Y),
+    kernel=gf.kernels.SquaredExponential(),
+)
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
+)
+print_summary(model)
 
 # %%
-# here's a plot of the raw data.
-fig, ax = plt.subplots(1, 1, figsize=(12, 6))
-_ = ax.plot(X, Y, "kx")
+mean, var = model.predict_y(X_plot)
+plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())
 
 # %% [markdown]
 # ### Data structure
 #
-# In this case, we need to let the model know which group each data point belongs to. We'll use a similar trick to the above, stacking the group identifier with the data:
+# To model the different noise groups, we need to let the model know which group each data point belongs to. We'll do that by appending a column with the group to the $x$ data:
 
 # %%
-Y_data = np.hstack([Y, groups])
+X_and_group = np.hstack([X, get_group(X)])
+X_plot_and_group = np.hstack([X_plot, get_group(X_plot)])
 
 # %% [markdown]
-# ### Build a likelihood
+# ### Use multiple functions for the noise variance
+#
+# To model this we will create two noise variance functions, each of which are just a constant, and then switch between them depending on the group labels.
 #
-# This time, we'll use a builtin likelihood, `SwitchedLikelihood`, which is a container for other likelihoods, and applies them to the first `Y_data` column depending on the index in the second. We're able to access and optimize the parameters of those likelihoods. Here, we'll (incorrectly) initialize the variances of our likelihoods to 1, to demonstrate how we can recover reasonable values for these through maximum-likelihood estimation.
+# Notice that we initialize the constant functions to a positive value. They would have defaulted to `0.0`, but that wouldn't work for a variance, which must be strictly positive.
 
 # %%
-likelihood = gpflow.likelihoods.SwitchedLikelihood(
-    [gpflow.likelihoods.Gaussian(variance=1.0), gpflow.likelihoods.Gaussian(variance=1.0)]
+model = gf.models.GPR(
+    data=(X_and_group, Y),
+    kernel=gf.kernels.SquaredExponential(active_dims=[0]),
+    likelihood=gf.likelihoods.Gaussian(
+        variance=gf.functions.SwitchedFunction(
+            [
+                gf.functions.Constant(1.0),
+                gf.functions.Constant(1.0),
+            ]
+        )
+    ),
 )
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
+)
+print_summary(model)
 
 # %%
-# model construction (notice that num_latent_gps is 1)
-kernel = gpflow.kernels.Matern52(lengthscales=0.5)
-model = gpflow.models.VGP((X, Y_data), kernel=kernel, likelihood=likelihood, num_latent_gps=1)
+mean, var = model.predict_y(X_plot_and_group)
+plot_distribution(X, Y, X_plot, mean.numpy(), var.numpy())
+
+# %% [markdown]
+# ## Demo 3: Empirical noise variance
+#
+# In this demo we will assume that you have multiple measurements at each $x$ location, and want to use the empirical variance at each location.
 
+# %% [markdown]
+# ### Generate data
 
 # %%
-for _ in range(reduce_in_tests(1000)):
-    natgrad.minimize(model.training_loss, [(model.q_mu, model.q_sqrt)])
+n_data = reduce_in_tests(20)
+n_repeats = reduce_in_tests(10)
+
+
+def generate_empiricial_noise_data() -> Tuple[gf.base.AnyNDArray, gf.base.AnyNDArray]:
+    rng = np.random.default_rng(42)  # for reproducibility
+    X = rng.uniform(0.0, 1.0, (n_data, 1))
+    signal = np.sin(6 * X)
+    noise_scale = rng.uniform(0.1, 1.0, (n_data, 1))
+    noise = noise_scale * rng.standard_normal((n_data, n_repeats))
+    Y = signal + noise
+    return X, Y
+
+
+X, Y = generate_empiricial_noise_data()
+Y_mean = np.mean(Y, axis=-1, keepdims=True)
+Y_var = np.var(Y, axis=-1, keepdims=True)
 
 # %% [markdown]
-# We've now fitted the VGP model to the data, but without optimizing over the hyperparameters. Plotting the data, we see that the fit is not terrible, but hasn't made use of our knowledge of the varying noise.
+# For the sake of plotting we'll create flat lists of (x, y) pairs:
 
 # %%
-# let's do some plotting!
-xx = np.linspace(-5, 5, 200)[:, None]
+X_flat = np.broadcast_to(X, Y.shape).reshape((-1, 1))
+Y_flat = Y.reshape((-1, 1))
+plot_distribution(X_flat, Y_flat)
 
-mu, var = model.predict_f(xx)
-
-fig, ax = plt.subplots(1, 1, figsize=(12, 6))
-ax.plot(xx, mu, "C0")
-ax.plot(xx, mu + 2 * np.sqrt(var), "C0", lw=0.5)
-ax.plot(xx, mu - 2 * np.sqrt(var), "C0", lw=0.5)
+# %% [markdown]
+# ### Data structure
+#
+# Our GPs don't like it when multiple data points occupy the same $x$ location. So we'll reduce this dataset to the means and the variances at each $x$, and use a custom function to model inject the empirical variances into the model.
 
-ax.plot(X, Y, "C1x", mew=2)
-_ = ax.set_xlim(-5, 5)
+# %%
+Y_mean = np.mean(Y, axis=-1, keepdims=True)
 
 # %% [markdown]
-# ### Optimizing the noise variances
-# Here we'll optimize over both the noise variance and the variational parameters, applying natural gradients interleaved with the Adam optimizer.
+# ### Fit a naive model
 #
-# As before, we mark the variational parameters as not trainable, in order to train them separately with the natural gradient method.
+# Again we'll start by fitting a naive GPR, to demonstrate the problem.
+#
+# Notice that this case is somewhat different from what we did above. Here we are not modelling $Y$, but the mean of $Y$, so the noise variance is not our uncertainty of $Y$, but our uncertainty of the mean of $Y$. Furthermore, while the naive model assumes a constant noise variance, the custom noise variance function we will develop below will only model the noise variance at the points where we have data, thus we cannot plot a continuous uncertainty. Instead, the shaded gray area is the confidence interval of $f$, and we use vertical black lines to indicate the confidence interval of the $Y$ mean.
 
 # %%
-likelihood = gpflow.likelihoods.SwitchedLikelihood(
-    [gpflow.likelihoods.Gaussian(variance=1.0), gpflow.likelihoods.Gaussian(variance=1.0)]
+model = gf.models.GPR(
+    data=(X, Y_mean),
+    kernel=gf.kernels.SquaredExponential(),
+)
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
 )
-kernel = gpflow.kernels.Matern52(lengthscales=0.5)
-model = gpflow.models.VGP((X, Y_data), kernel=kernel, likelihood=likelihood, num_latent_gps=1)
+print_summary(model)
 
-set_trainable(model.q_mu, False)
-set_trainable(model.q_sqrt, False)
+# %%
+f_mean, f_var = model.predict_f(X_plot)
+y_mean_mean, y_mean_var = model.predict_y(X)
+
+plot_distribution(
+    X_flat,
+    Y_flat,
+    X_plot,
+    f_mean.numpy(),
+    f_var.numpy(),
+    X,
+    y_mean_mean.numpy(),
+    y_mean_var.numpy(),
+)
 
-for _ in range(reduce_in_tests(1000)):
-    natgrad.minimize(model.training_loss, [(model.q_mu, model.q_sqrt)])
-    adam.minimize(model.training_loss, model.trainable_variables)
 
 # %% [markdown]
-# ### Plotting the fitted model
+# ### Create custom function for the noise variance
 #
-# Now that the noise variances have been estimated, we can see the final model fit.
-# The predictive variance is higher on the left side of the plot, where we know that the data have different variance.
-# We'll plot the known underlying function in green to see how effectively we've recovered the ground truth.
-# We can also print the model to examine the estimated noise variances:
+# We're modelling the `Y_mean` and its standard error is `Y_var / n_repeats`. We will create a simple function that ignores its input and returns these values as a constant. This will obviously only work for the `X` that corresponds to the `Y` we computed the variance of - that is good enough for model training, but notice that it will not allow us to use `predict_y`.
 
 # %%
-# let's do some plotting!
-xx = np.linspace(-5, 5, 200)[:, None]
+class FixedVarianceOfMean(gf.functions.Function):
+    def __init__(self, Y: gf.base.AnyNDArray):
+        n_repeats = Y.shape[-1]
+        self.var_mean = np.var(Y, axis=-1, keepdims=True) / n_repeats
 
-mu, var = model.predict_f(xx)
+    @inherit_check_shapes
+    def __call__(self, X: gf.base.TensorType) -> tf.Tensor:
+        return self.var_mean
 
-fig, ax = plt.subplots(1, 1, figsize=(12, 6))
-ax.plot(xx, mu, "C0")
-ax.plot(xx, mu + 2 * np.sqrt(var), "C0", lw=0.5)
-ax.plot(xx, mu - 2 * np.sqrt(var), "C0", lw=0.5)
 
-ax.plot(X, Y, "C1x", mew=2)
-ax.set_xlim(-5, 5)
-_ = ax.plot(xx, 2.5 * np.sin(6 * xx) + np.cos(3 * xx), "C2--")
+# %% [markdown]
+# Now we can plug that into our model:
 
 # %%
-from gpflow.utilities import print_summary
+model = gf.models.GPR(
+    data=(X, Y_mean),
+    kernel=gf.kernels.SquaredExponential(active_dims=[0]),
+    likelihood=gf.likelihoods.Gaussian(variance=FixedVarianceOfMean(Y)),
+)
+gf.optimizers.Scipy().minimize(
+    model.training_loss, model.trainable_variables, options=optimizer_config
+)
+print_summary(model)
 
-print_summary(model, fmt="notebook")
+# %%
+f_mean, f_var = model.predict_f(X_plot)
+y_mean_mean, y_mean_var = model.predict_y(X)
+
+plot_distribution(
+    X_flat,
+    Y_flat,
+    X_plot,
+    f_mean.numpy(),
+    f_var.numpy(),
+    X,
+    y_mean_mean.numpy(),
+    y_mean_var.numpy(),
+)
+
+# %% [markdown]
+# You may notice that the data points fall outside our confidence interval. Again, this is because the confidence interval is for the mean of $Y$, and not the $Y$ points themselves.
 
 # %% [markdown]
 # ## Further reading
diff --git a/gpflow/functions.py b/gpflow/functions.py
--- a/gpflow/functions.py
+++ b/gpflow/functions.py
@@ -287,11 +287,11 @@ def __call__(self, X: TensorType) -> tf.Tensor:
         )  # X = X[:,:-1]
 
         # split up X into chunks corresponding to the relevant likelihoods
-        x_list = tf.dynamic_partition(X, ind, len(self.meanfunctions))
+        x_list = tf.dynamic_partition(X, ind, len(self.functions))
         # apply the likelihood-function to each section of the data
-        results = [m(x) for x, m in zip(x_list, self.meanfunctions)]
+        results = [m(x) for x, m in zip(x_list, self.functions)]
         # stitch the results back together
-        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.meanfunctions))
+        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.functions))
         return tf.dynamic_stitch(partitions, results)
 
 


Test patch: