collapsed_elbo#

gpjax.objectives.collapsed_elbo(variational_family, data)[source]#

Compute a single step of the collapsed evidence lower bound.

Compute the evidence lower bound under this model. In short, this requires evaluating the expectation of the model’s log-likelihood under the variational approximation. To this, we sum the KL divergence from the variational posterior to the prior. This collapsed bound is evaluated on the full dataset supplied in data and does not apply minibatch scaling.

Example

>>> import gpjax as gpx
>>> import jax.numpy as jnp

>>> xtrain = jnp.linspace(0, 1).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)

>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood

>>> z = jnp.linspace(0, 1, 10).reshape(-1, 1)
>>> q = gpx.variational_families.CollapsedVariationalGaussian(
...     posterior=posterior, inducing_inputs=z
... )

>>> gpx.objectives.collapsed_elbo(q, D)
Parameters:

  • variational_family (VF) – The variational approximation for whose parameters we should maximise the ELBO with respect to.

  • data (Dataset) – The training data for which we should maximise the ELBO with respect to.

Returns:

The evidence lower bound of the variational approximation.

Return type:

ScalarFloat