# Copyright 2018 The TensorFlow Probability Authors.
#
# 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.
# ============================================================================
"""The VariationalGaussianProcess distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import gaussian_process
from tensorflow_probability.python.distributions import independent
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import mvn_linear_operator
from tensorflow_probability.python.distributions import normal
from tensorflow_probability.python.internal import dtype_util

__all__ = [
    'VariationalGaussianProcess',
]


def _add_diagonal_shift(matrix, shift):
  return tf.linalg.set_diag(
      matrix, tf.linalg.diag_part(matrix) + shift, name='add_diagonal_shift')


def _solve_cholesky_factored_system(
    cholesky_factor, rhs, name=None):
  with tf.name_scope(
      name or '_solve_cholesky_factored_system') as scope:
    cholesky_factor = tf.convert_to_tensor(value=cholesky_factor,
                                           name='cholesky_factor')
    rhs = tf.convert_to_tensor(value=rhs, name='rhs')
    lin_op = tf.linalg.LinearOperatorLowerTriangular(
        cholesky_factor, name=scope)
    return lin_op.solve(lin_op.solve(rhs), adjoint=True)


def _solve_cholesky_factored_system_vec(cholesky_factor, rhs, name=None):
  with tf.name_scope(
      name or '_solve_cholesky_factored_system') as scope:
    cholesky_factor = tf.convert_to_tensor(
        value=cholesky_factor, name='cholesky_factor')
    rhs = tf.convert_to_tensor(value=rhs, name='rhs')
    lin_op = tf.linalg.LinearOperatorLowerTriangular(
        cholesky_factor, name=scope)
    return lin_op.solvevec(lin_op.solvevec(rhs), adjoint=True)


class VariationalGaussianProcess(
    mvn_linear_operator.MultivariateNormalLinearOperator):
  """Posterior predictive of a variational Gaussian process.

  This distribution implements the variational Gaussian process (VGP), as
  described in [Titsias, 2009][1] and [Hensman, 2013][2]. The VGP is an
  inducing point-based approximation of an exact GP posterior
  (see Mathematical Details, below). Ultimately, this Distribution class
  represents a marginal distrbution over function values at a
  collection of `index_points`. It is parameterized by

    - a kernel function,
    - a mean function,
    - the (scalar) observation noise variance of the normal likelihood,
    - a set of index points,
    - a set of inducing index points, and
    - the parameters of the (full-rank, Gaussian) variational posterior
      distribution over function values at the inducing points, conditional on
      some observations.

  A VGP is "trained" by selecting any kernel parameters, the locations of the
  inducing index points, and the variational parameters. [Titsias, 2009][1] and
  [Hensman, 2013][2] describe a variational lower bound on the marginal log
  likelihood of observed data, which this class offers through the
  `variational_loss` method (this is the negative lower bound, for convenience
  when plugging into a TF Optimizer's `minimize` function).
  Training may be done in minibatches.

  [Titsias, 2009][1] describes a closed form for the optimal variational
  parameters, in the case of sufficiently small observational data (ie,
  small enough to fit in memory but big enough to warrant approximating the GP
  posterior). A method to compute these optimal parameters in terms of the full
  observational data set is provided as a staticmethod,
  `optimal_variational_posterior`. It returns a
  `MultivariateNormalLinearOperator` instance with optimal location and
  scale parameters.

  #### Mathematical Details

  ##### Notation

  We will in general be concerned about three collections of index points, and
  it'll be good to give them names:

    * `x[1], ..., x[N]`: observation index points -- locations of our observed
      data.
    * `z[1], ..., z[M]`: inducing index points  -- locations of the
      "summarizing" inducing points
    * `t[1], ..., t[P]`: predictive index points -- locations where we are
      making posterior predictions based on observations and the variational
      parameters.

  To lighten notation, we'll use `X, Z, T` to denote the above collections.
  Similarly, we'll denote by `f(X)` the collection of function values at each of
  the `x[i]`, and by `Y`, the collection of (noisy) observed data at each `x[i].
  We'll denote kernel matrices generated from pairs of index points as `K_tt`,
  `K_xt`, `K_tz`, etc, e.g.,

  ```none
           | k(t[1], z[1])    k(t[1], z[2])  ...  k(t[1], z[M]) |
    K_tz = | k(t[2], z[1])    k(t[2], z[2])  ...  k(t[2], z[M]) |
           |      ...              ...                 ...      |
           | k(t[P], z[1])    k(t[P], z[2])  ...  k(t[P], z[M]) |
  ```

  ##### Preliminaries

  A Gaussian process is an indexed collection of random variables, any finite
  collection of which are jointly Gaussian. Typically, the index set is some
  finite-dimensional, real vector space, and indeed we make this assumption in
  what follows. The GP may then be thought of as a distribution over functions
  on the index set. Samples from the GP are functions *on the whole index set*;
  these can't be represented in finite compute memory, so one typically works
  with the marginals at a finite collection of index points. The properties of
  the GP are entirely determined by its mean function `m` and covariance
  function `k`. The generative process, assuming a mean-zero normal likelihood
  with stddev `sigma`, is

  ```none
    f ~ GP(m, k)

    Y | f(X) ~ Normal(f(X), sigma),   i = 1, ... , N
  ```

  In finite terms (ie, marginalizing out all but a finite number of f(X)'sigma),
  we can write

  ```none
    f(X) ~ MVN(loc=m(X), cov=K_xx)

    Y | f(X) ~ Normal(f(X), sigma),   i = 1, ... , N
  ```

  Posterior inference is possible in analytical closed form but becomes
  intractible as data sizes get large. See [Rasmussen, 2006][3] for details.

  ##### The VGP

  The VGP is an inducing point-based approximation of an exact GP posterior,
  where two approximating assumptions have been made:

    1. function values at non-inducing points are mutually independent
       conditioned on function values at the inducing points,
    2. the (expensive) posterior over function values at inducing points
       conditional on obseravtions is replaced with an arbitrary (learnable)
       full-rank Gaussian distribution,

       ```none
         q(f(Z)) = MVN(loc=m, scale=S),
       ```

       where `m` and `S` are parameters to be chosen by optimizing an evidence
       lower bound (ELBO).

  The posterior predictive distribution becomes

  ```none
    q(f(T)) = integral df(Z) p(f(T) | f(Z)) q(f(Z))
            = MVN(loc = A @ m, scale = B^(1/2))
  ```

  where

  ```none
    A = K_tz @ K_zz^-1
    B = K_tt - A @ (K_zz - S S^T) A^T
  ```

  ***The approximate posterior predictive distribution `q(f(T))` is what the
  `VariationalGaussianProcess` class represents.***

  Model selection in this framework entails choosing the kernel parameters,
  inducing point locations, and variational parameters. We do this by optimizing
  a variational lower bound on the marginal log likelihood of observed data. The
  lower bound takes the following form (see [Titsias, 2009][1] and
  [Hensman, 2013][2] for details on the derivation):

  ```none
    L(Z, m, S, Y) = (
        MVN(loc=(K_zx @ K_zz^-1) @ m, scale_diag=sigma).log_prob(Y) -
        (Tr(K_xx - K_zx @ K_zz^-1 @ K_xz) +
         Tr(S @ S^T @ K_zz^1 @ K_zx @ K_xz @ K_zz^-1)) / (2 * sigma^2) -
        KL(q(f(Z)) || p(f(Z))))
  ```

  where in the final KL term, `p(f(Z))` is the GP prior on inducing point
  function values. This variational lower bound can be computed on minibatches
  of the full data set `(X, Y)`. A method to compute the *negative* variational
  lower bound is implemented as `VariationalGaussianProcess.variational_loss`.

  ##### Optimal variational parameters

  As described in [Titsias, 2009][1], a closed form optimum for the variational
  location and scale parameters, `m` and `S`, can be computed when the
  observational data are not prohibitively voluminous. The
  `optimal_variational_posterior` function to computes the optimal variational
  posterior distribution over inducing point function values in terms of the GP
  parameters (mean and kernel functions), inducing point locations, observation
  index points, and observations. Note that the inducing index point locations
  must still be optimized even when these parameters are known functions of the
  inducing index points. The optimal parameters are computed as follows:

  ```none
    C = sigma^-2 (K_zz + K_zx @ K_xz)^-1

    optimal Gaussian covariance: K_zz @ C @ K_zz
    optimal Gaussian location: sigma^-2 K_zz @ C @ K_zx @ Y
  ```

  #### Usage Examples

  Here's an example of defining and training a VariationalGaussianProcess on
  some toy generated data.

  ```python
  # We'll use double precision throughout for better numerics.
  dtype = np.float64

  # Generate noisy data from a known function.
  f = lambda x: np.exp(-x[..., 0]**2 / 20.) * np.sin(1. * x[..., 0])
  true_observation_noise_variance_ = dtype(1e-1) ** 2

  num_training_points_ = 100
  x_train_ = np.stack(
      [np.random.uniform(-6., 0., [num_training_points_/ 2 , 1]),
       np.random.uniform(1., 10., [num_training_points_/ 2 , 1])],
      axis=0).astype(dtype)
  y_train_ = (f(x_train_) +
              np.random.normal(
                  0., np.sqrt(true_observation_noise_variance_),
                  [num_training_points_]).astype(dtype))

  # Create kernel with trainable parameters, and trainable observation noise
  # variance variable. Each of these is constrained to be positive.
  amplitude = (tf.nn.softplus(tf.Variable(-1., dtype=dtype, name='amplitude')))
  length_scale = (1e-5 +
                  tf.nn.softplus(
                      tf.Variable(-3., dtype=dtype, name='length_scale')))
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)

  observation_noise_variance = tf.nn.softplus(
      tf.Variable(0, dtype=dtype, name='observation_noise_variance'))

  # Create trainable inducing point locations and variational parameters.
  num_inducing_points_ = 20

  inducing_index_points = tf.Variable(
      initial_inducing_points_, dtype=dtype,
      name='inducing_index_points')
  variational_inducing_observations_loc = tf.Variable(
      np.zeros([num_inducing_points_], dtype=dtype),
      name='variational_inducing_observations_loc')
  variational_inducing_observations_scale = tf.Variable(
      np.eye(num_inducing_points_, dtype=dtype),
      name='variational_inducing_observations_scale')

  # These are the index point locations over which we'll construct the
  # (approximate) posterior predictive distribution.
  num_predictive_index_points_ = 500
  index_points_ = np.linspace(-13, 13,
                              num_predictive_index_points_,
                              dtype=dtype)[..., np.newaxis]


  # Construct our variational GP Distribution instance.
  vgp = tfd.VariationalGaussianProcess(
      kernel,
      index_points=index_points_,
      inducing_index_points=inducing_index_points,
      variational_inducing_observations_loc=variational_inducing_observations_loc,
      variational_inducing_observations_scale=variational_inducing_observations_scale,
      observation_noise_variance=observation_noise_variance)

  # For training, we use some simplistic numpy-based minibatching.
  batch_size = 64
  x_train_batch = tf.placeholder(dtype, [batch_size, 1], name='x_train_batch')
  y_train_batch = tf.placeholder(dtype, [batch_size], name='y_train_batch')

  # Create the loss function we want to optimize.
  loss = vgp.variational_loss(
      observations=y_train_batch,
      observation_index_points=x_train_batch,
      kl_weight=float(batch_size) / float(num_training_points_))

  optimizer = tf.train.AdamOptimizer()
  train_op = optimizer.minimize(loss)

  num_iters = 10000
  num_logs = 10
  with tf.Session() as sess:
    for i in range(num_iters):
      batch_idxs = np.random.randint(num_training_points_, size=[batch_size])
      x_train_batch_ = x_train_[batch_idxs, ...]
      y_train_batch_ = y_train_[batch_idxs]

      [_, loss_] = sess.run([train_op, loss],
                            feed_dict={x_train_batch: x_train_batch_,
                                       y_train_batch: y_train_batch_})
      if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
        print(i, loss_)

  # Generate a plot with
  #   - the posterior predictive mean
  #   - training data
  #   - inducing index points (plotted vertically at the mean of the variational
  #     posterior over inducing point function values)
  #   - 50 posterior predictive samples

  num_samples = 50
  [
      samples_,
      mean_,
      inducing_index_points_,
      variational_loc_,
  ] = sess.run([
      vgp.sample(num_samples),
      vgp.mean(),
      inducing_index_points,
      variational_inducing_observations_loc
  ])
  plt.figure(figsize=(15, 5))
  plt.scatter(inducing_index_points_[..., 0], variational_loc_
              marker='x', s=50, color='k', zorder=10)
  plt.scatter(x_train_[..., 0], y_train_, color='#00ff00', zorder=9)
  plt.plot(np.tile(index_points_[..., 0], num_samples),
           samples_.T, color='r', alpha=.1)
  plt.plot(index_points_, mean_, color='k')
  plt.plot(index_points_, f(index_points_), color='b')
  ```

  # Here we use the same data setup, but compute the optimal variational
  # parameters instead of training them.
  ```python
  # We'll use double precision throughout for better numerics.
  dtype = np.float64

  # Generate noisy data from a known function.
  f = lambda x: np.exp(-x[..., 0]**2 / 20.) * np.sin(1. * x[..., 0])
  true_observation_noise_variance_ = dtype(1e-1) ** 2

  num_training_points_ = 1000
  x_train_ = np.random.uniform(-10., 10., [num_training_points_, 1])
  y_train_ = (f(x_train_) +
              np.random.normal(
                  0., np.sqrt(true_observation_noise_variance_),
                  [num_training_points_]))

  # Create kernel with trainable parameters, and trainable observation noise
  # variance variable. Each of these is constrained to be positive.
  amplitude = (tf.nn.softplus(
    tf.Variable(.54, dtype=dtype, name='amplitude', use_resource=True)))
  length_scale = (
    1e-5 +
    tf.nn.softplus(
      tf.Variable(.54, dtype=dtype, name='length_scale', use_resource=True)))
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)

  observation_noise_variance = tf.nn.softplus(
      tf.Variable(
        .54, dtype=dtype, name='observation_noise_variance', use_resource=True))

  # Create trainable inducing point locations and variational parameters.
  num_inducing_points_ = 10

  inducing_index_points = tf.Variable(
      np.linspace(-10., 10., num_inducing_points_)[..., np.newaxis],
      dtype=dtype, name='inducing_index_points', use_resource=True)

  variational_loc, variational_scale = (
      tfd.VariationalGaussianProcess.optimal_variational_posterior(
          kernel=kernel,
          inducing_index_points=inducing_index_points,
          observation_index_points=x_train_,
          observations=y_train_,
          observation_noise_variance=observation_noise_variance))

  # These are the index point locations over which we'll construct the
  # (approximate) posterior predictive distribution.
  num_predictive_index_points_ = 500
  index_points_ = np.linspace(-13, 13,
                              num_predictive_index_points_,
                              dtype=dtype)[..., np.newaxis]

  # Construct our variational GP Distribution instance.
  vgp = tfd.VariationalGaussianProcess(
      kernel,
      index_points=index_points_,
      inducing_index_points=inducing_index_points,
      variational_inducing_observations_loc=variational_loc,
      variational_inducing_observations_scale=variational_scale,
      observation_noise_variance=observation_noise_variance)

  # For training, we use some simplistic numpy-based minibatching.
  batch_size = 64
  x_train_batch = tf.placeholder(dtype, [batch_size, 1], name='x_train_batch')
  y_train_batch = tf.placeholder(dtype, [batch_size], name='y_train_batch')

  # Create the loss function we want to optimize.
  loss = vgp.variational_loss(
      observations=y_train_batch,
      observation_index_points=x_train_batch,
      kl_weight=float(batch_size) / float(num_training_points_))

  optimizer = tf.train.AdamOptimizer(learning_rate=.01)
  train_op = optimizer.minimize(loss)

  num_iters = 300
  num_logs = 10
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(num_iters):
      batch_idxs = np.random.randint(num_training_points_, size=[batch_size])
      x_train_batch_ = x_train_[batch_idxs, ...]
      y_train_batch_ = y_train_[batch_idxs]

      [_, loss_] = sess.run([train_op, loss],
                            feed_dict={x_train_batch: x_train_batch_,
                                       y_train_batch: y_train_batch_})
      if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
        print(i, loss_)

    # Generate a plot with
    #   - the posterior predictive mean
    #   - training data
    #   - inducing index points (plotted vertically at the mean of the
    #     variational posterior over inducing point function values)
    #   - 50 posterior predictive samples

    num_samples = 50
    [
        samples_,
        mean_,
        inducing_index_points_,
        variational_loc_,
    ] = sess.run([
        vgp.sample(num_samples),
        vgp.mean(),
        inducing_index_points,
        variational_loc
    ])
    plt.figure(figsize=(15, 5))
    plt.scatter(inducing_index_points_[..., 0], variational_loc_,
                marker='x', s=50, color='k', zorder=10)
    plt.scatter(x_train_[..., 0], y_train_, color='#00ff00', alpha=.1, zorder=9)
    plt.plot(np.tile(index_points_, num_samples),
             samples_.T, color='r', alpha=.1)
    plt.plot(index_points_, mean_, color='k')
    plt.plot(index_points_, f(index_points_), color='b')

  ```

  #### References

  [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
       Regression", 2009.
       http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf
  [2]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
       https://arxiv.org/abs/1309.6835
  [3]: Carl Rasmussen, Chris Williams. Gaussian Processes For Machine Learning,
       2006. http://www.gaussianprocess.org/gpml/
  """

  def __init__(self,
               kernel,
               index_points,
               inducing_index_points,
               variational_inducing_observations_loc,
               variational_inducing_observations_scale,
               mean_fn=None,
               observation_noise_variance=0.,
               predictive_noise_variance=0.,
               jitter=1e-6,
               validate_args=False,
               allow_nan_stats=False,
               name='VariataionalGaussianProcess'):
    """Instantiate a VariationalGaussianProcess Distribution.

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the VGP is defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below).
        Ultimately the VariationalGaussianProcess distribution corresponds to an
        `e1`-dimensional multivariate normal. The batch shape must be
        broadcastable with `kernel.batch_shape`, the batch shape of
        `inducing_index_points`, and any batch dims yielded by `mean_fn`.
      inducing_index_points: `float` `Tensor` of locations of inducing points in
        the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
        like `index_points`. The batch shape components needn't be identical to
        those of `index_points`, but must be broadcast compatible with them.
      variational_inducing_observations_loc: `float` `Tensor`; the mean of the
        (full-rank Gaussian) variational posterior over function values at the
        inducing points, conditional on observed data. Shape has the form `[b1,
        ..., bB, e2]`, where `b1, ..., bB` is broadcast compatible with other
        parameters' batch shapes, and `e2` is the number of inducing points.
      variational_inducing_observations_scale: `float` `Tensor`; the scale
        matrix of the (full-rank Gaussian) variational posterior over function
        values at the inducing points, conditional on observed data. Shape has
        the form `[b1, ..., bB, e2, e2]`, where `b1, ..., bB` is broadcast
        compatible with other parameters and `e2` is the number of inducing
        points.
      mean_fn: Python `callable` that acts on index points to produce a (batch
        of) vector(s) of mean values at those index points. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      predictive_noise_variance: `float` `Tensor` representing additional
        variance in the posterior predictive model. If `None`, we simply re-use
        `observation_noise_variance` for the posterior predictive noise. If set
        explicitly, however, we use the given value. This allows us, for
        example, to omit predictive noise variance (by setting this to zero) to
        obtain noiseless posterior predictions of function values, conditioned
        on noisy observations.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
        Default value: `False`.
      allow_nan_stats: Python `bool`, default `True`. When `True`,
        statistics (e.g., mean, mode, variance) use the value "`NaN`" to
        indicate the result is undefined. When `False`, an exception is raised
        if one or more of the statistic's batch members are undefined.
        Default value: `False`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "VariationalGaussianProcess".

    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.
    """
    parameters = dict(locals())
    with tf.name_scope(
        name or 'VariationalGaussianProcess') as name:
      dtype = dtype_util.common_dtype(
          [kernel,
           index_points,
           inducing_index_points,
           variational_inducing_observations_loc,
           variational_inducing_observations_scale,
           observation_noise_variance,
           predictive_noise_variance,
           jitter], tf.float32)

      index_points = tf.convert_to_tensor(
          value=index_points, dtype=dtype, name='index_points')
      inducing_index_points = tf.convert_to_tensor(
          value=inducing_index_points, dtype=dtype,
          name='inducing_index_points')
      variational_inducing_observations_loc = tf.convert_to_tensor(
          value=variational_inducing_observations_loc, dtype=dtype,
          name='variational_inducing_observations_loc')
      variational_inducing_observations_scale = tf.convert_to_tensor(
          value=variational_inducing_observations_scale, dtype=dtype,
          name='variational_inducing_observations_scale')
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')
      if predictive_noise_variance is None:
        predictive_noise_variance = observation_noise_variance
      else:
        predictive_noise_variance = tf.convert_to_tensor(
            value=predictive_noise_variance, dtype=dtype,
            name='predictive_noise_variance')
      jitter = tf.convert_to_tensor(
          value=jitter, dtype=dtype, name='jitter')

      self._kernel = kernel
      self._index_points = index_points
      self._inducing_index_points = inducing_index_points
      self._variational_inducing_observations_posterior = (
          mvn_linear_operator.MultivariateNormalLinearOperator(
              loc=variational_inducing_observations_loc,
              scale=tf.linalg.LinearOperatorFullMatrix(
                  variational_inducing_observations_scale),
              name='variational_inducing_observations_posterior'))

      # Default to a constant zero function, borrowing the dtype from
      # index_points to ensure consistency.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')
      self._mean_fn = mean_fn
      self._observation_noise_variance = observation_noise_variance
      self._predictive_noise_variance = predictive_noise_variance
      self._jitter = jitter

      with tf.name_scope('init'):
        # We let t and z denote predictive and inducing index points, resp.
        kzz = _add_diagonal_shift(
            kernel.matrix(inducing_index_points, inducing_index_points),
            jitter)

        self._chol_kzz = tf.linalg.cholesky(kzz)
        self._kzz_inv_varloc = _solve_cholesky_factored_system_vec(
            self._chol_kzz,
            (variational_inducing_observations_loc -
             mean_fn(inducing_index_points)),
            name='kzz_inv_varloc')

        loc, scale = self._compute_posterior_predictive_params()

        super(VariationalGaussianProcess, self).__init__(
            loc=loc,
            scale=scale,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            name=name)
        self._parameters = parameters
        self._graph_parents = [
            index_points,
            inducing_index_points,
            variational_inducing_observations_loc,
            variational_inducing_observations_scale,
            observation_noise_variance,
            jitter]

  def _compute_posterior_predictive_params(self):
    ktt = _add_diagonal_shift(
        self._kernel.matrix(self._index_points, self._index_points),
        self._jitter)
    kzt = tf.linalg.LinearOperatorFullMatrix(
        self._kernel.matrix(self._inducing_index_points, self._index_points))

    kzz_inv_kzt = tf.linalg.LinearOperatorFullMatrix(
        _solve_cholesky_factored_system(
            self._chol_kzz, kzt.to_dense(), name='kzz_inv_kzt'))

    var_cov = tf.linalg.LinearOperatorFullMatrix(
        self._variational_inducing_observations_posterior.covariance())
    posterior_predictive_cov = (
        ktt -
        kzt.matmul(kzz_inv_kzt.to_dense(), adjoint=True) +
        kzz_inv_kzt.matmul(var_cov.matmul(kzz_inv_kzt.to_dense()),
                           adjoint=True))

    # Add predictive_noise_variance
    posterior_predictive_cov = _add_diagonal_shift(
        posterior_predictive_cov, self._predictive_noise_variance)

    scale = tf.linalg.LinearOperatorLowerTriangular(
        tf.linalg.cholesky(posterior_predictive_cov))

    loc = (self._mean_fn(self._index_points) +
           kzt.matvec(self._kzz_inv_varloc, adjoint=True))

    return loc, scale

  @property
  def mean_fn(self):
    return self._mean_fn

  @property
  def kernel(self):
    return self._kernel

  @property
  def index_points(self):
    return self._index_points

  @property
  def inducing_index_points(self):
    return self._inducing_index_points

  @property
  def variational_inducing_observations_loc(self):
    return self._variational_inducing_observations_posterior.loc

  @property
  def variational_inducing_observations_scale(self):
    return self._variational_inducing_observations_posterior.scale

  @property
  def observation_noise_variance(self):
    return self._observation_noise_variance

  @property
  def predictive_noise_variance(self):
    return self._predictive_noise_variance

  @property
  def jitter(self):
    return self._jitter

  def _covariance(self):
    return self._covariance_matrix

  def variational_loss(self,
                       observations,
                       observation_index_points=None,
                       kl_weight=1.,
                       name='variational_loss'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(name or 'variational_gp_loss'):
      if observation_index_points is None:
        observation_index_points = self._index_points
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=self._dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=self._dtype, name='observations')
      kl_weight = tf.convert_to_tensor(
          value=kl_weight, dtype=self._dtype,
          name='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      kzx = self.kernel.matrix(self._inducing_index_points,
                               observation_index_points)

      kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)
      loc = (self._mean_fn(observation_index_points) +
             kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))

      likelihood = independent.Independent(
          normal.Normal(
              loc=loc,
              scale=tf.sqrt(self._observation_noise_variance + self._jitter),
              name='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      obs_ll = likelihood.log_prob(observations)

      chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)
      chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)
      kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)

      kxx_diag = tf.linalg.diag_part(
          self.kernel.matrix(
              observation_index_points, observation_index_points))
      ktilde_trace_term = (
          tf.reduce_sum(input_tensor=kxx_diag, axis=-1) -
          tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      other_trace_term = tf.reduce_sum(
          input_tensor=(
              self._variational_inducing_observations_posterior.scale.matmul(
                  kzz_inv_kzx) ** 2),
          axis=[-2, -1])

      trace_term = (.5 * (ktilde_trace_term + other_trace_term) /
                    self._observation_noise_variance)

      inducing_prior = gaussian_process.GaussianProcess(
          kernel=self._kernel,
          mean_fn=self._mean_fn,
          index_points=self._inducing_index_points,
          observation_noise_variance=self._observation_noise_variance)

      kl_term = kl_weight * kullback_leibler.kl_divergence(
          self._variational_inducing_observations_posterior,
          inducing_prior)

      lower_bound = (obs_ll - trace_term - kl_term)

      return -tf.reduce_mean(input_tensor=lower_bound)

  @staticmethod
  def optimal_variational_posterior(
      kernel,
      inducing_index_points,
      observation_index_points,
      observations,
      observation_noise_variance,
      mean_fn=None,
      jitter=1e-6,
      name=None):
    """Model selection for optimal variational hyperparameters.

    Given the full training set (parameterized by `observations` and
    `observation_index_points`), compute the optimal variational
    location and scale for the VGP. This is based of the method suggested
    in [Titsias, 2009][1].

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      inducing_index_points: `float` `Tensor` of locations of inducing points in
        the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
        like `observation_index_points`. The batch shape components needn't be
        identical to those of `observation_index_points`, but must be broadcast
        compatible with them.
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below).
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      mean_fn: Python `callable` that acts on index points to produce a (batch
        of) vector(s) of mean values at those index points. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "optimal_variational_posterior".
    Returns:
      loc, scale: Tuple representing the variational location and scale.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
         Regression", 2009.
         http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf
    """

    with tf.name_scope(name or 'optimal_variational_posterior'):
      dtype = dtype_util.common_dtype(
          [inducing_index_points,
           observation_index_points,
           observations,
           observation_noise_variance,
           jitter], tf.float32)

      inducing_index_points = tf.convert_to_tensor(
          value=inducing_index_points,
          dtype=dtype, name='inducing_index_points')
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=dtype, name='observations')
      observation_noise_variance = tf.convert_to_tensor(
          value=observation_noise_variance,
          dtype=dtype,
          name='observation_noise_variance')
      jitter = tf.convert_to_tensor(
          value=jitter, dtype=dtype, name='jitter')

      # Default to a constant zero function.
      if mean_fn is None:
        mean_fn = lambda x: tf.zeros([1], dtype=dtype)
      else:
        if not callable(mean_fn):
          raise ValueError('`mean_fn` must be a Python callable')

      # z are the inducing points and x are the observation index points.
      kzz = kernel.matrix(inducing_index_points, inducing_index_points)
      kzx = kernel.matrix(inducing_index_points, observation_index_points)

      noise_var_inv = tf.math.reciprocal(observation_noise_variance)

      sigma_inv = _add_diagonal_shift(
          kzz + noise_var_inv * tf.matmul(kzx, kzx, adjoint_b=True),
          jitter)

      chol_sigma_inv = tf.linalg.cholesky(sigma_inv)

      kzx_lin_op = tf.linalg.LinearOperatorFullMatrix(kzx)
      kzx_obs = kzx_lin_op.matvec(
          observations - mean_fn(observation_index_points))
      kzz_lin_op = tf.linalg.LinearOperatorFullMatrix(kzz)
      loc = (mean_fn(inducing_index_points) +
             noise_var_inv * kzz_lin_op.matvec(
                 _solve_cholesky_factored_system_vec(chol_sigma_inv, kzx_obs)))

      chol_sigma_inv_lin_op = tf.linalg.LinearOperatorLowerTriangular(
          chol_sigma_inv)
      scale = chol_sigma_inv_lin_op.solve(kzz)

      return loc, scale
