class_docstr: 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 observations 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
  import matplotlib.pyplot as plt
  import numpy as np
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  tfb = tfp.bijectors
  tfd = tfp.distributions
  tfk = tfp.math.psd_kernels
  
  # 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.concatenate(
      [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 = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='amplitude')
  length_scale = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='length_scale')
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)
  
  observation_noise_variance = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='observation_noise_variance')
  
  # Create trainable inducing point locations and variational parameters.
  num_inducing_points_ = 20
  inducing_index_points = tf.Variable(
      np.linspace(-5., 5., num_inducing_points_)[..., np.newaxis],
      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
  
  optimizer = tf_keras.optimizers.Adam(learning_rate=.1)
  
  @tf.function
  def optimize(x_train_batch, y_train_batch):
    with tf.GradientTape() as tape:
      # 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_))
    grads = tape.gradient(loss, vgp.trainable_variables)
    optimizer.apply_gradients(zip(grads, vgp.trainable_variables))
    return loss
  
  num_iters = 10000
  num_logs = 10
  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 = optimize(x_train_batch, y_train_batch)
  
    if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
      print(i, loss.numpy())
  
  # 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 = vgp.sample(num_samples).numpy()
  mean = vgp.mean().numpy()
  inducing_index_points_ = inducing_index_points.numpy()
  variational_loc = variational_inducing_observations_loc.numpy()
  
  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_, (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 = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='amplitude')
  length_scale = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='length_scale')
  kernel = tfk.ExponentiatedQuadratic(
      amplitude=amplitude,
      length_scale=length_scale)
  
  observation_noise_variance = tfp.util.TransformedVariable(
      1., tfb.Softplus(), dtype=dtype, name='observation_noise_variance')
  
  # 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')
  
  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
  
  optimizer = tf_keras.optimizers.Adam(learning_rate=.05, beta_1=.5, beta_2=.99)
  
  @tf.function
  def optimize(x_train_batch, y_train_batch):
    with tf.GradientTape() as tape:
      # 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_))
    grads = tape.gradient(loss, vgp.trainable_variables)
    optimizer.apply_gradients(zip(grads, vgp.trainable_variables))
    return loss
  
  num_iters = 300
  num_logs = 10
  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 = optimize(x_train_batch, y_train_batch)
    if i % (num_iters / num_logs) == 0 or i + 1 == num_iters:
      print(i, loss.numpy())
  
  # 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_ = vgp.sample(num_samples).numpy()
  mean_ = vgp.mean().numpy()
  inducing_index_points_ = inducing_index_points.numpy()
  variational_loc_ = variational_loc.numpy()
  
  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/
  [4]: Hensman, J., Matthews, A. G., Filippone M., Ghahramani Z. "MCMC for
       Variationally Sparse Gaussian Processes"
       https://arxiv.org/abs/1506.04000
class_docstr_tok: ["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","ca","n'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","observations","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","import","matplotlib.pyplot","as","plt","import","numpy","as","np","import","tensorflow.compat.v2","as","tf","import","tensorflow_probability","as","tfp","tfb","=","tfp.bijectors","tfd","=","tfp.distributions","tfk","=","tfp.math.psd_kernels","#","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.concatenate","(","[","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","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='amplitude","'",")","length_scale","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='length_scale","'",")","kernel","=","tfk.ExponentiatedQuadratic","(","amplitude=amplitude",",","length_scale=length_scale",")","observation_noise_variance","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='observation_noise_variance","'",")","#","Create","trainable","inducing","point","locations","and","variational","parameters",".","num_inducing_points_","=","20","inducing_index_points","=","tf.Variable","(","np.linspace","(","-5.",",","5.",",","num_inducing_points_",")","[","...",",","np.newaxis","]",",","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","optimizer","=","tf_keras.optimizers.Adam","(","learning_rate=.1",")","@","tf.function","def","optimize","(","x_train_batch",",","y_train_batch",")",":","with","tf.GradientTape","(",")","as","tape",":","#","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_",")",")","grads","=","tape.gradient","(","loss",",","vgp.trainable_variables",")","optimizer.apply_gradients","(","zip","(","grads",",","vgp.trainable_variables",")",")","return","loss","num_iters","=","10000","num_logs","=","10","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","=","optimize","(","x_train_batch",",","y_train_batch",")","if","i","%","(","num_iters","\/","num_logs",")","==","0","or","i","+","1","==","num_iters",":","print","(","i",",","loss.numpy","(",")",")","#","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","=","vgp.sample","(","num_samples",")",".numpy","(",")","mean","=","vgp.mean","(",")",".numpy","(",")","inducing_index_points_","=","inducing_index_points.numpy","(",")","variational_loc","=","variational_inducing_observations_loc.numpy","(",")","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_",",","(","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","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='amplitude","'",")","length_scale","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='length_scale","'",")","kernel","=","tfk.ExponentiatedQuadratic","(","amplitude=amplitude",",","length_scale=length_scale",")","observation_noise_variance","=","tfp.util.TransformedVariable","(","1.",",","tfb.Softplus","(",")",",","dtype=dtype",",","name='observation_noise_variance","'",")","#","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","'",")","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","optimizer","=","tf_keras.optimizers.Adam","(","learning_rate=.05",",","beta_1=.5",",","beta_2=.99",")","@","tf.function","def","optimize","(","x_train_batch",",","y_train_batch",")",":","with","tf.GradientTape","(",")","as","tape",":","#","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_",")",")","grads","=","tape.gradient","(","loss",",","vgp.trainable_variables",")","optimizer.apply_gradients","(","zip","(","grads",",","vgp.trainable_variables",")",")","return","loss","num_iters","=","300","num_logs","=","10","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","=","optimize","(","x_train_batch",",","y_train_batch",")","if","i","%","(","num_iters","\/","num_logs",")","==","0","or","i","+","1","==","num_iters",":","print","(","i",",","loss.numpy","(",")",")","#","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_","=","vgp.sample","(","num_samples",")",".numpy","(",")","mean_","=","vgp.mean","(",")",".numpy","(",")","inducing_index_points_","=","inducing_index_points.numpy","(",")","variational_loc_","=","variational_loc.numpy","(",")","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\/","[","4","]",":","Hensman",",","J.",",","Matthews",",","A.","G.",",","Filippone","M.",",","Ghahramani","Z",".","``","MCMC","for","Variationally","Sparse","Gaussian","Processes","''","https",":","\/\/arxiv.org\/abs\/1506.04000"]
docstr: 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, e, f1, ..., fF]` and returns a `Tensor` whose shape
      is (broadcastable with) `[b1, ..., bB, e]`.
      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.
    cholesky_fn: Callable which takes a single (batch) matrix argument and
      returns a Cholesky-like lower triangular factor.  Default value: `None`,
      in which case `make_cholesky_with_jitter_fn` is used with the `jitter`
      parameter.
    use_whitening_transform: Python `bool`. Whether to reparameterize
      the variational scale as `m = chol(K_zz)m', s = chol(K_zz)s'`, where
      `chol(K_zz)` is the cholesky factor of the kernel at the
      `inducing_index_points`, `m` is the variational location, `s` is the
      variational scale, and `m', s'` is the reparameterization when this
      parameter is True.
      Default value: `False`.
    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.
docstr_tok: ["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","need","n'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",",","e",",","f1",",","...",",","fF","]","`","and","returns","a","`","Tensor","`","whose","shape","is","(","broadcastable","with",")","`","[","b1",",","...",",","bB",",","e","]","`",".","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",".","cholesky_fn",":","Callable","which","takes","a","single","(","batch",")","matrix","argument","and","returns","a","Cholesky-like","lower","triangular","factor",".","Default","value",":","`","None","`",",","in","which","case","`","make_cholesky_with_jitter_fn","`","is","used","with","the","`","jitter","`","parameter",".","use_whitening_transform",":","Python","`","bool","`",".","Whether","to","reparameterize","the","variational","scale","as","`","m","=","chol","(","K_zz",")","m","'",",","s","=","chol","(","K_zz",")","s","'","`",",","where","`","chol","(","K_zz",")","`","is","the","cholesky","factor","of","the","kernel","at","the","`","inducing_index_points","`",",","`","m","`","is","the","variational","location",",","`","s","`","is","the","variational","scale",",","and","`","m","'",",","s","'","`","is","the","reparameterization","when","this","parameter","is","True",".","Default","value",":","`","False","`",".","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","."]
code: def __init__(
      self,
      kernel,
      index_points,
      inducing_index_points,
      variational_inducing_observations_loc,
      variational_inducing_observations_scale,
      mean_fn=None,
      observation_noise_variance=None,
      predictive_noise_variance=None,
      cholesky_fn=None,
      use_whitening_transform=False,
      jitter=1e-6,
      validate_args=False,
      allow_nan_stats=False,
      name="VariationalGaussianProcess",
  ):
      """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, e, f1, ..., fF]` and returns a `Tensor` whose shape
          is (broadcastable with) `[b1, ..., bB, e]`.
          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.
        cholesky_fn: Callable which takes a single (batch) matrix argument and
          returns a Cholesky-like lower triangular factor.  Default value: `None`,
          in which case `make_cholesky_with_jitter_fn` is used with the `jitter`
          parameter.
        use_whitening_transform: Python `bool`. Whether to reparameterize
          the variational scale as `m = chol(K_zz)m', s = chol(K_zz)s'`, where
          `chol(K_zz)` is the cholesky factor of the kernel at the
          `inducing_index_points`, `m` is the variational location, `s` is the
          variational scale, and `m', s'` is the reparameterization when this
          parameter is True.
          Default value: `False`.
        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 = tensor_util.convert_nonref_to_tensor(
              index_points, dtype=dtype, name="index_points"
          )
          inducing_index_points = tensor_util.convert_nonref_to_tensor(
              inducing_index_points,
              dtype=dtype,
              name="inducing_index_points",
          )
          variational_inducing_observations_loc = (
              tensor_util.convert_nonref_to_tensor(
                  variational_inducing_observations_loc,
                  dtype=dtype,
                  name="variational_inducing_observations_loc",
              )
          )
          variational_inducing_observations_scale = (
              tensor_util.convert_nonref_to_tensor(
                  variational_inducing_observations_scale,
                  dtype=dtype,
                  name="variational_inducing_observations_scale",
              )
          )
          observation_noise_variance = (
              tensor_util.convert_nonref_to_tensor(
                  observation_noise_variance,
                  dtype=dtype,
                  name="observation_noise_variance",
              )
          )
  
          predictive_noise_variance = (
              tensor_util.convert_nonref_to_tensor(
                  predictive_noise_variance,
                  dtype=dtype,
                  name="predictive_noise_variance",
              )
          )
          if predictive_noise_variance is None:
              predictive_noise_variance = observation_noise_variance
  
          self._use_whitening_transform = use_whitening_transform
          self._kernel = kernel
          self._index_points = index_points
          self._inducing_index_points = inducing_index_points
          self._variational_inducing_observations_loc = (
              variational_inducing_observations_loc
          )
          self._variational_inducing_observations_scale = (
              variational_inducing_observations_scale
          )
          self._variational_inducing_observations_posterior = (
              mvn_linear_operator.MultivariateNormalLinearOperator(
                  loc=variational_inducing_observations_loc,
                  scale=tf.linalg.LinearOperatorLowerTriangular(
                      variational_inducing_observations_scale
                  ),
                  name="variational_inducing_observations_posterior",
              )
          )
  
          is_zero_mean = False
          # 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)
              is_zero_mean = True
          else:
              if not callable(mean_fn):
                  raise ValueError(
                      "`mean_fn` must be a Python callable"
                  )
          self._prior_mean_fn = mean_fn
  
          variational_kernel = _VariationalKernel(
              kernel,
              inducing_index_points,
              variational_inducing_observations_scale,
              cholesky_fn=cholesky_fn,
              use_whitening_transform=use_whitening_transform,
              jitter=jitter,
          )
  
          # Set up the prior over function values at inducing points using the
          # original kernel and mean_fn (as opposed to the variational posterior
          # predictive kernel and mean fn we define below). This is a distribution
          # over latent function values so there is no observation noise term.
          if use_whitening_transform:
              # Under the whitening transformation, this is a standard multivariate
              # normal (possibly shifted).
              mean = None
              if not is_zero_mean:
  
                  def solve_mean(z):
                      chol_kzz = (
                          tf.linalg.LinearOperatorLowerTriangular(
                              variational_kernel.chol_kzz()
                          )
                      )
                      # mean_fn(z) needs to have the same dimensions as chol_kzz. We do
                      # this by an explicit broadcast.
                      mean = mean_fn(z) + tf.zeros_like(
                          variational_inducing_observations_loc
                      )
                      return chol_kzz.solvevec(mean)
  
                  mean = tfp_util.DeferredTensor(
                      inducing_index_points,
                      transform_fn=solve_mean,
                      shape=None,
                      name="inducing_prior_mean",
                  )
  
              self._inducing_prior = (
                  mvn_linear_operator.MultivariateNormalLinearOperator(
                      loc=mean,
                      scale=tf.linalg.LinearOperatorIdentity(
                          variational_inducing_observations_scale.shape[
                              -1
                          ],
                          dtype=dtype,
                      ),
                  )
              )
          else:
              self._inducing_prior = gaussian_process.GaussianProcess(
                  kernel=kernel,
                  mean_fn=mean_fn,
                  index_points=inducing_index_points,
              )
  
          self._predictive_noise_variance = predictive_noise_variance
          self._vgp_observation_noise_variance = (
              observation_noise_variance
          )
  
          posterior_predictive_mean_fn = (
              _make_posterior_predictive_mean_fn(
                  kernel=kernel,
                  mean_fn=mean_fn,
                  inducing_index_points=inducing_index_points,
                  variational_inducing_observations_loc=(
                      variational_inducing_observations_loc
                  ),
                  use_whitening_transform=use_whitening_transform,
                  chol_kzz_fn=variational_kernel.chol_kzz,
              )
          )
  
          super(VariationalGaussianProcess, self).__init__(
              kernel=variational_kernel,
              mean_fn=posterior_predictive_mean_fn,
              index_points=index_points,
              cholesky_fn=cholesky_fn,
              jitter=jitter,
              # What the GP super class calls "observation noise variance" we call
              # here the "predictive noise variance". We use the observation noise
              # variance for the fit/solve process, and predictive for downstream
              # computations like sampling.
              observation_noise_variance=predictive_noise_variance,
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              name=name,
          )
          self._parameters = parameters
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python import util
  from tensorflow_probability.python.bijectors import fill_scale_tril
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import cholesky_util
  from tensorflow_probability.python.distributions import distribution
  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.distributions.internal import stochastic_process_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import linalg
  from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel
  from tensorflow_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 1 to create a new VariationalGaussianProcess object from inherited base classes: gaussian_process, distribution with example: obj = VariationalGaussianProcess(kernel, index_points, inducing_index_points, variational_inducing_observations_loc, variational_inducing_observations_scale, mean_fn, observation_noise_variance, predictive_noise_variance, cholesky_fn, use_whitening_transform, jitter, validate_args, allow_nan_stats, name)