docstr: Tests that the factorised KL terms sum up to the true KL.
docstr_tok: ["Tests","that","the","factorised","KL","terms","sum","up","to","the","true","KL","."]
code: 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
          )
  
file_path: trfl/trfl/distribution_ops_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import itertools
  from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow_probability
  from trfl import distribution_ops
label_desc: Use image node_id 2 for calling the FactorisedKLGaussianTest obj's underlying member method code with example usage: obj.testFactorisedKLGaussian(dist1_type, dist2_type) without return types