def testFactorisedKLGaussian(self, dist1_type, dist2_type):
    """Tests that the factorised KL terms sum up to the true KL."""
    dist1, dist1_mean, dist1_cov = self._create_gaussian(dist1_type)
    dist2, dist2_mean, dist2_cov = self._create_gaussian(dist2_type)
    both_diagonal = _is_diagonal(dist1.scale) and _is_diagonal(dist2.scale)
    if both_diagonal:
      dist1_cov = dist1.parameters['scale_diag']
      dist2_cov = dist2.parameters['scale_diag']
    kl = tfp.distributions.kl_divergence(dist1, dist2)
    kl_mean, kl_cov = distribution_ops.factorised_kl_gaussian(
        dist1_mean,
        dist1_cov,
        dist2_mean,
        dist2_cov,
        both_diagonal=both_diagonal)
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      actual_kl, kl_mean_np, kl_cov_np = sess.run([kl, kl_mean, kl_cov])
      self.assertAllClose(actual_kl, kl_mean_np + kl_cov_np, rtol=1e-4)
Tests that the factorised KL terms sum up to the true KL.
trfl/distribution_ops_test.py