Add priors to kernel hyperparameters to loss

gpflux/layers/gp_layer.py

@@ -284,7 +284,10 @@ def call(self, inputs: TensorType, *args: List[Any], **kwargs: Dict[str, Any]) -
         outputs = super().call(inputs, *args, **kwargs)
 
         if kwargs.get("training"):
-            loss_per_datapoint = self.prior_kl() / self.num_data
+            log_prior = tf.add_n([p.log_prior_density() for p in self.kernel.trainable_parameters])
+            loss = self.prior_kl() - log_prior
+            loss_per_datapoint = loss / self.num_data
+
         else:
             # TF quirk: add_loss must always add a tensor to compile
             loss_per_datapoint = tf.constant(0.0, dtype=default_float())

setup.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 from pathlib import Path
-from setuptools import find_namespace_packages, setup
 
+from setuptools import find_namespace_packages, setup
 
 requirements = [
     "deprecated",

tests/integration/test_svgp_equivalence.py

@@ -19,14 +19,28 @@
 import numpy as np
 import pytest
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 import gpflow
+from gpflow import Parameter
+from gpflow.models.model import RegressionData
+from gpflow.utilities import positive, to_default_float
 
 import gpflux
 
 tf.keras.backend.set_floatx("float64")
 
 
+class LogPrior_ELBO_SVGP(gpflow.models.SVGP):
+    """
+    SVGP model that takes into account the log_prior in the ELBO
+    """
+
+    def elbo(self, data: RegressionData) -> tf.Tensor:
+        loss_prior = tf.add_n([p.log_prior_density() for p in self.trainable_parameters])
+        return super().elbo(data) + loss_prior
+
+
 def load_data():
     path = os.path.join(os.path.dirname(__file__), "..", "snelson1d.npz")
     data = np.load(path)
@@ -48,6 +62,9 @@ def make_dataset(data, as_dict=True):
 
 def make_kernel_likelihood_iv():
     kernel = gpflow.kernels.SquaredExponential(variance=0.7, lengthscales=0.6)
+    kernel.lengthscales.prior = tfp.distributions.LogNormal(
+        to_default_float(1.0), to_default_float(0.5)
+    )
     likelihood = gpflow.likelihoods.Gaussian(variance=0.08)
     Z = np.linspace(0, 6, 20)[:, np.newaxis]
     inducing_variable = gpflow.inducing_variables.InducingPoints(Z)
@@ -56,7 +73,7 @@ def make_kernel_likelihood_iv():
 
 
 def create_gpflow_svgp(kernel, likelihood, inducing_variable):
-    return gpflow.models.SVGP(kernel, likelihood, inducing_variable)
+    return LogPrior_ELBO_SVGP(kernel, likelihood, inducing_variable)
 
 
 def create_gp_layer(kernel, inducing_variable, num_data):