Generalize implementation of Gaussian expectations of Gaussian kernels by j-wilson · Pull Request #1607

Generalize implementation of Gaussian expectations of Gaussian kernels by j-wilson · Pull Request #1607 · GPflow/GPflow

Proposed changes
Extend the existing implementation of Gaussian expectations of Gaussian kernels $E_{p(x)}[k1(Z1, x) k2(x, Z2)]$ to support different kernels and different inducing locations.

import tensorflow as tf
from gpflow.kernels import SquaredExponential
from gpflow.expectations import expectation
from gpflow.inducing_variables import InducingPoints
from gpflow.probability_distributions import Gaussian

N = 4  # num. uncertain test locations
M = 3  # num. inducing locations per kernel
D = 2  # dimensionality of input domain

mx = tf.random.uniform([N, D],  dtype=tf.float64)
Lxx = tf.random.normal([N, D, D],  dtype=tf.float64)
Sxx = tf.matmul(Lxx, Lxx, transpose_b=True)
px = Gaussian(mx, Sxx)

k1 = SquaredExponential(lengthscales=0.01 + tf.random.uniform([D], dtype=tf.float64))
Z1 = InducingPoints(tf.random.uniform([M, D], dtype=tf.float64))

k2 = SquaredExponential(lengthscales=0.01 + tf.random.uniform([D], dtype=tf.float64))
Z2 = InducingPoints(tf.random.uniform([M, D], dtype=tf.float64))

eK = expectation(px, (k1, Z1), (k2, Z2))  # [N, M, M]