class_docstr: Well tested because this is a simple override of the base class.
class_docstr_tok: ["Well","tested","because","this","is","a","simple","override","of","the","base","class","."]
code: def testSample(self):
      # TODO(jvdillon): This test should be the basis of a new test fixture which
      # is applied to every distribution. When we make this fixture, we'll also
      # separate the analytical- and sample-based tests as well as for each
      # function tested. For now, we group things so we can recycle one batch of
      # samples (thus saving resources).
  
      mu = np.array([-1.0, 1, 0.5], dtype=np.float32)
      diag_large = np.array([1.0, 0.5, 0.75], dtype=np.float32)
      diag_small = np.array([-1.1, 1.2], dtype=np.float32)
      v = np.array(
          [[0.7, 0.8], [0.9, 1], [0.5, 0.6]], dtype=np.float32
      )  # shape: [k, r] = [3, 2]
  
      true_mean = mu
      true_scale = np.diag(diag_large) + np.matmul(
          np.matmul(v, np.diag(diag_small)), v.T
      )
      true_covariance = np.matmul(true_scale, true_scale.T)
      true_variance = np.diag(true_covariance)
      true_stddev = np.sqrt(true_variance)
  
      dist = mvdpl.MultivariateNormalDiagPlusLowRank(
          loc=mu,
          scale_diag=diag_large,
          scale_perturb_factor=v,
          scale_perturb_diag=diag_small,
          validate_args=True,
      )
  
      # The following distributions will test the KL divergence calculation.
      mvn_identity = mvn_diag.MultivariateNormalDiag(
          loc=np.array([1.0, 2, 0.25], dtype=np.float32),
          validate_args=True,
      )
      mvn_scaled = mvn_diag.MultivariateNormalDiag(
          loc=mvn_identity.loc,
          scale_diag=2.2 * tf.ones_like(mvn_identity.loc),
          validate_args=True,
      )
      mvd = mvn_diag.MultivariateNormalDiag(
          loc=mvn_identity.loc,
          scale_diag=np.array([0.5, 1.5, 1.0], dtype=np.float32),
          validate_args=True,
      )
      mvn_chol = mvn_tril.MultivariateNormalTriL(
          loc=np.array([1.0, 2, -1], dtype=np.float32),
          scale_tril=np.array(
              [[6.0, 0, 0], [2, 5, 0], [1, 3, 4]], dtype=np.float32
          )
          / 10.0,
          validate_args=True,
      )
  
      scale = dist.scale.to_dense()
  
      n = int(70e3)
      samps = dist.sample(n, seed=test_util.test_seed())
      sample_mean = tf.reduce_mean(samps, axis=0)
      x = samps - sample_mean
      sample_covariance = tf.matmul(x, x, transpose_a=True) / n
  
      sample_kl_identity = tf.reduce_mean(
          dist.log_prob(samps) - mvn_identity.log_prob(samps), axis=0
      )
      analytical_kl_identity = kullback_leibler.kl_divergence(
          dist, mvn_identity
      )
  
      sample_kl_scaled = tf.reduce_mean(
          dist.log_prob(samps) - mvn_scaled.log_prob(samps), axis=0
      )
      analytical_kl_scaled = kullback_leibler.kl_divergence(
          dist, mvn_scaled
      )
  
      sample_kl_diag = tf.reduce_mean(
          dist.log_prob(samps) - mvd.log_prob(samps), axis=0
      )
      analytical_kl_diag = kullback_leibler.kl_divergence(dist, mvd)
  
      sample_kl_chol = tf.reduce_mean(
          dist.log_prob(samps) - mvn_chol.log_prob(samps), axis=0
      )
      analytical_kl_chol = kullback_leibler.kl_divergence(
          dist, mvn_chol
      )
  
      n = int(10e3)
      baseline = mvn_diag.MultivariateNormalDiag(
          loc=np.array([-1.0, 0.25, 1.25], dtype=np.float32),
          scale_diag=np.array([1.5, 0.5, 1.0], dtype=np.float32),
          validate_args=True,
      )
      samps = baseline.sample(n, seed=test_util.test_seed())
  
      sample_kl_identity_diag_baseline = tf.reduce_mean(
          baseline.log_prob(samps) - mvn_identity.log_prob(samps),
          axis=0,
      )
      analytical_kl_identity_diag_baseline = (
          kullback_leibler.kl_divergence(baseline, mvn_identity)
      )
  
      sample_kl_scaled_diag_baseline = tf.reduce_mean(
          baseline.log_prob(samps) - mvn_scaled.log_prob(samps), axis=0
      )
      analytical_kl_scaled_diag_baseline = (
          kullback_leibler.kl_divergence(baseline, mvn_scaled)
      )
  
      sample_kl_diag_diag_baseline = tf.reduce_mean(
          baseline.log_prob(samps) - mvd.log_prob(samps), axis=0
      )
      analytical_kl_diag_diag_baseline = kullback_leibler.kl_divergence(
          baseline, mvd
      )
  
      sample_kl_chol_diag_baseline = tf.reduce_mean(
          baseline.log_prob(samps) - mvn_chol.log_prob(samps), axis=0
      )
      analytical_kl_chol_diag_baseline = kullback_leibler.kl_divergence(
          baseline, mvn_chol
      )
  
      [
          sample_mean_,
          analytical_mean_,
          sample_covariance_,
          analytical_covariance_,
          analytical_variance_,
          analytical_stddev_,
          scale_,
          sample_kl_identity_,
          analytical_kl_identity_,
          sample_kl_scaled_,
          analytical_kl_scaled_,
          sample_kl_diag_,
          analytical_kl_diag_,
          sample_kl_chol_,
          analytical_kl_chol_,
          sample_kl_identity_diag_baseline_,
          analytical_kl_identity_diag_baseline_,
          sample_kl_scaled_diag_baseline_,
          analytical_kl_scaled_diag_baseline_,
          sample_kl_diag_diag_baseline_,
          analytical_kl_diag_diag_baseline_,
          sample_kl_chol_diag_baseline_,
          analytical_kl_chol_diag_baseline_,
      ] = self.evaluate(
          [
              sample_mean,
              dist.mean(),
              sample_covariance,
              dist.covariance(),
              dist.variance(),
              dist.stddev(),
              scale,
              sample_kl_identity,
              analytical_kl_identity,
              sample_kl_scaled,
              analytical_kl_scaled,
              sample_kl_diag,
              analytical_kl_diag,
              sample_kl_chol,
              analytical_kl_chol,
              sample_kl_identity_diag_baseline,
              analytical_kl_identity_diag_baseline,
              sample_kl_scaled_diag_baseline,
              analytical_kl_scaled_diag_baseline,
              sample_kl_diag_diag_baseline,
              analytical_kl_diag_diag_baseline,
              sample_kl_chol_diag_baseline,
              analytical_kl_chol_diag_baseline,
          ]
      )
  
      sample_variance_ = np.diag(sample_covariance_)
      sample_stddev_ = np.sqrt(sample_variance_)
  
      tf1.logging.vlog(2, "true_mean:\n{}  ".format(true_mean))
      tf1.logging.vlog(2, "sample_mean:\n{}".format(sample_mean_))
      tf1.logging.vlog(
          2, "analytical_mean:\n{}".format(analytical_mean_)
      )
  
      tf1.logging.vlog(
          2, "true_covariance:\n{}".format(true_covariance)
      )
      tf1.logging.vlog(
          2, "sample_covariance:\n{}".format(sample_covariance_)
      )
      tf1.logging.vlog(
          2, "analytical_covariance:\n{}".format(analytical_covariance_)
      )
  
      tf1.logging.vlog(2, "true_variance:\n{}".format(true_variance))
      tf1.logging.vlog(
          2, "sample_variance:\n{}".format(sample_variance_)
      )
      tf1.logging.vlog(
          2, "analytical_variance:\n{}".format(analytical_variance_)
      )
  
      tf1.logging.vlog(2, "true_stddev:\n{}".format(true_stddev))
      tf1.logging.vlog(2, "sample_stddev:\n{}".format(sample_stddev_))
      tf1.logging.vlog(
          2, "analytical_stddev:\n{}".format(analytical_stddev_)
      )
  
      tf1.logging.vlog(2, "true_scale:\n{}".format(true_scale))
      tf1.logging.vlog(2, "scale:\n{}".format(scale_))
  
      tf1.logging.vlog(
          2,
          "kl_identity:  analytical:{}  sample:{}".format(
              analytical_kl_identity_, sample_kl_identity_
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_scaled:    analytical:{}  sample:{}".format(
              analytical_kl_scaled_, sample_kl_scaled_
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_diag:      analytical:{}  sample:{}".format(
              analytical_kl_diag_, sample_kl_diag_
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_chol:      analytical:{}  sample:{}".format(
              analytical_kl_chol_, sample_kl_chol_
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_identity_diag_baseline:  analytical:{}  sample:{}".format(
              analytical_kl_identity_diag_baseline_,
              sample_kl_identity_diag_baseline_,
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_scaled_diag_baseline:  analytical:{}  sample:{}".format(
              analytical_kl_scaled_diag_baseline_,
              sample_kl_scaled_diag_baseline_,
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_diag_diag_baseline:  analytical:{}  sample:{}".format(
              analytical_kl_diag_diag_baseline_,
              sample_kl_diag_diag_baseline_,
          ),
      )
  
      tf1.logging.vlog(
          2,
          "kl_chol_diag_baseline:  analytical:{}  sample:{}".format(
              analytical_kl_chol_diag_baseline_,
              sample_kl_chol_diag_baseline_,
          ),
      )
  
      self.assertAllClose(true_mean, sample_mean_, atol=0.0, rtol=0.05)
      self.assertAllClose(
          true_mean, analytical_mean_, atol=0.0, rtol=1e-6
      )
  
      self.assertAllClose(
          true_covariance, sample_covariance_, atol=0.0, rtol=0.06
      )
      self.assertAllClose(
          true_covariance, analytical_covariance_, atol=0.0, rtol=1e-6
      )
  
      self.assertAllClose(
          true_variance, sample_variance_, atol=0.0, rtol=0.05
      )
      self.assertAllClose(
          true_variance, analytical_variance_, atol=0.0, rtol=1e-6
      )
  
      self.assertAllClose(
          true_stddev, sample_stddev_, atol=0.0, rtol=0.05
      )
      self.assertAllClose(
          true_stddev, analytical_stddev_, atol=0.0, rtol=1e-6
      )
  
      self.assertAllClose(true_scale, scale_, atol=0.0, rtol=1e-6)
  
      self.assertAllClose(
          sample_kl_identity_,
          analytical_kl_identity_,
          atol=0.0,
          rtol=0.05,
      )
      self.assertAllClose(
          sample_kl_scaled_, analytical_kl_scaled_, atol=0.0, rtol=0.05
      )
      self.assertAllClose(
          sample_kl_diag_, analytical_kl_diag_, atol=0.0, rtol=0.05
      )
      self.assertAllClose(
          sample_kl_chol_, analytical_kl_chol_, atol=0.0, rtol=0.05
      )
  
      self.assertAllClose(
          sample_kl_identity_diag_baseline_,
          analytical_kl_identity_diag_baseline_,
          atol=0.0,
          rtol=0.05,
      )
      self.assertAllClose(
          sample_kl_scaled_diag_baseline_,
          analytical_kl_scaled_diag_baseline_,
          atol=0.0,
          rtol=0.05,
      )
      self.assertAllClose(
          sample_kl_diag_diag_baseline_,
          analytical_kl_diag_diag_baseline_,
          atol=0.0,
          rtol=0.05,
      )
      self.assertAllClose(
          sample_kl_chol_diag_baseline_,
          analytical_kl_chol_diag_baseline_,
          atol=0.0,
          rtol=0.05,
      )
  
file_path: probability/tensorflow_probability/python/distributions/mvn_diag_plus_low_rank_test.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_diag_plus_low_rank
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 7 for calling the MultivariateNormalDiagPlusLowRankTest obj's underlying member method code with example usage: obj.testSample() without return types