def factorised_kl_gaussian(dist1_mean,
                           dist1_covariance_or_scale,
                           dist2_mean,
                           dist2_covariance_or_scale,
                           both_diagonal=False):
  """Compute the KL divergence KL(dist1, dist2) between two Gaussians.

  The KL is factorised into two terms - `kl_mean` and `kl_cov`. This
  factorisation is specific to multivariate gaussian distributions and arises
  from its analytic form.
  Specifically, if we assume two multivariate Gaussian distributions with rank
  k and means, M1 and M2 and variance S1 and S2, the analytic KL can be written
  out as:

  D_KL(N0 || N1) = 0.5 * (tr(inv(S1) * S0) + ln(det(S1)/det(S0)) - k +
                         (M1 - M0).T * inv(S1) * (M1 - M0))

  The terms on the first row correspond to the covariance factor and the terms
  on the second row correspond to the mean factor in the factorized KL.
  These terms can thus be used to independently control how much the mean and
  covariance between the two gaussians can vary.

  This implementation ensures that gradient flow is equivalent to calling
  `tfp.distributions.kl_divergence` once.

  More details on the equation can be found here:
  https://stats.stackexchange.com/questions/60680/kl-divergence-between-two-multivariate-gaussians


  Args:
    dist1_mean: The mean of the first Multivariate Gaussian distribution.
    dist1_covariance_or_scale: The covariance or scale of the first Multivariate
      Gaussian distribution. In cases where *both* distributions are Gaussians
      with diagonal covariance matrices (for instance, if both are instances of
      `tfp.distributions.MultivariateNormalDiag`), then the `scale` can be
      passed in instead and the `both_diagonal` flag must be set to `True`.
      A more efficient sparse computation path is used in this case. For all
      other cases, the full covariance matrix must be passed in.
    dist2_mean: The mean of the second Multivariate Gaussian distribution.
    dist2_covariance_or_scale: The covariance or scale tensor of the second
      Multivariate Gaussian distribution, as for `dist1_covariance_or_scale`.
    both_diagonal: A `bool` indicating that both dist1 and dist2 are diagonal
      matrices. A more efficient sparse computation is used in this case.

  Returns:
    A tuple consisting of (`kl_mean`, `kl_cov`) which correspond to the mean and
    the covariance factorisation of the KL.
  """
  if both_diagonal:
    dist1_mean_rank = dist1_mean.get_shape().ndims
    dist1_covariance_or_scale.get_shape().assert_has_rank(dist1_mean_rank)
    dist2_mean_rank = dist2_mean.get_shape().ndims
    dist2_covariance_or_scale.get_shape().assert_has_rank(dist2_mean_rank)

    dist_type = tfp.distributions.MultivariateNormalDiag
  else:
    dist_type = tfp.distributions.MultivariateNormalFullCovariance

  # Recreate the distributions but with stop gradients on the mean and cov.
  dist1_stop_grad_mean = dist_type(
      tf.stop_gradient(dist1_mean), dist1_covariance_or_scale)
  dist2 = dist_type(dist2_mean, dist2_covariance_or_scale)

  # Now create a third distribution with the mean of dist1 and the variance of
  # dist2 and appropriate stop_gradients.
  dist3 = dist_type(dist1_mean, dist2_covariance_or_scale)
  dist3_stop_grad_mean = dist_type(
      tf.stop_gradient(dist1_mean), dist2_covariance_or_scale)

  # Finally get the two components of the KL between dist1 and dist2
  # using dist3
  kl_mean = tfp.distributions.kl_divergence(dist3, dist2)
  kl_cov = tfp.distributions.kl_divergence(dist1_stop_grad_mean,
                                           dist3_stop_grad_mean)
  return kl_mean, kl_cov
Compute the KL divergence KL(dist1, dist2) between two Gaussians.

  The KL is factorised into two terms - `kl_mean` and `kl_cov`. This
  factorisation is specific to multivariate gaussian distributions and arises
  from its analytic form.
  Specifically, if we assume two multivariate Gaussian distributions with rank
  k and means, M1 and M2 and variance S1 and S2, the analytic KL can be written
  out as:

  D_KL(N0 || N1) = 0.5 * (tr(inv(S1) * S0) + ln(det(S1)/det(S0)) - k +
                         (M1 - M0).T * inv(S1) * (M1 - M0))

  The terms on the first row correspond to the covariance factor and the terms
  on the second row correspond to the mean factor in the factorized KL.
  These terms can thus be used to independently control how much the mean and
  covariance between the two gaussians can vary.

  This implementation ensures that gradient flow is equivalent to calling
  `tfp.distributions.kl_divergence` once.

  More details on the equation can be found here:
  https://stats.stackexchange.com/questions/60680/kl-divergence-between-two-multivariate-gaussians


  Args:
    dist1_mean: The mean of the first Multivariate Gaussian distribution.
    dist1_covariance_or_scale: The covariance or scale of the first Multivariate
      Gaussian distribution. In cases where *both* distributions are Gaussians
      with diagonal covariance matrices (for instance, if both are instances of
      `tfp.distributions.MultivariateNormalDiag`), then the `scale` can be
      passed in instead and the `both_diagonal` flag must be set to `True`.
      A more efficient sparse computation path is used in this case. For all
      other cases, the full covariance matrix must be passed in.
    dist2_mean: The mean of the second Multivariate Gaussian distribution.
    dist2_covariance_or_scale: The covariance or scale tensor of the second
      Multivariate Gaussian distribution, as for `dist1_covariance_or_scale`.
    both_diagonal: A `bool` indicating that both dist1 and dist2 are diagonal
      matrices. A more efficient sparse computation is used in this case.

  Returns:
    A tuple consisting of (`kl_mean`, `kl_cov`) which correspond to the mean and
    the covariance factorisation of the KL.
  
trfl/distribution_ops.py