code: def testLogProbMatchesTransformedDistribution(self, dtype):
      # In this test, we start with a distribution supported on N x N SPD matrices
      # which factorizes as an LKJ-supported correlation matrix and a diagonal
      # of exponential random variables. The total number of independent
      # parameters is N(N-1)/2 + N = N(N+1)/2. Using the CholeskyOuterProduct
      # bijector (which requires N(N+1)/2 independent parameters), we map it to
      # the space of lower triangular Cholesky factors. We show that the resulting
      # distribution factorizes as the product of CholeskyLKJ and Rayleigh
      # distributions.
  
      # Given a sample from the 'LKJ + Exponential' distribution on SPD matrices,
      # and the corresponding log prob, returns the transformed sample and log
      # prob in Cholesky-space; which is furthermore factored into a diagonal
      # matrix and a Cholesky factor of a correlation matrix.
      def _get_transformed_sample_and_log_prob(
          lkj_exp_sample, lkj_exp_log_prob, dimension
      ):
          # We change variables to the space of SPD matrices parameterized by the
          # lower triangular entries (including the diagonal) of \Sigma.
          # The transformation is given by \Sigma = \sqrt{S} P \sqrt{S}.
          #
          # The jacobian matrix J of the forward transform has the block form
          # [[I, 0]; [*, D]]; where I is the N x N identity matrix; and D is an
          # N*(N - 1) square diagonal matrix and * need not be computed.
          # Here, N is the dimension.  D_ij (0<=i<j<n) equals:
          #  d/d(P_ij) \Sigma_ij = d/d(P_ij) \sqrt{S_i S_j} P_ij = \sqrt{S_i S_j}.
          # Hence, detJ = \prod_ij (i < j) \sqrt{S_i S_j}  [N(N-1)/2 terms]
          #             = \prod_i S_i^{.5 * (N - 1)}
          # [1] Zhenxun Wang and Yunan Wu and Haitao Chu
          # 'On equivalence of the LKJ distribution and the restricted Wishart
          # distribution'. 2018
          exp_variance, lkj_corr = lkj_exp_sample
          sqrt_exp_variance = tf.math.sqrt(
              exp_variance[..., tf.newaxis]
          )
          sigma_sample = (
              tf.linalg.matrix_transpose(sqrt_exp_variance)
              * lkj_corr
              * sqrt_exp_variance
          )
          sigma_log_prob = lkj_exp_log_prob - 0.5 * (
              dimension - 1
          ) * tf.reduce_sum(
              tf.math.log(tf.linalg.diag_part(sigma_sample)), axis=-1
          )
  
          # We change variables again to lower triangular L; where LL^T = \Sigma.
          # This is just inverse of the tfb.CholeskyOuterProduct bijector.
          cholesky_sigma_sample = tf.linalg.cholesky(sigma_sample)
          cholesky_sigma_log_prob = sigma_log_prob + invert.Invert(
              cholesky_outer_product.CholeskyOuterProduct()
          ).inverse_log_det_jacobian(
              cholesky_sigma_sample, event_ndims=2
          )
  
          # Change of variables to R, A; where L = RA; R is diagonal matrix
          # with each dimension's standard deviation and A is a Cholesky factor of a
          # correlation matrix. A is parameterized by its strictly lower triangular
          # entries; i.e. N(N-1)/2 entries.
          #
          # The jacobian determinant is the product of each row's transformation, as
          # each row is transformed independently as R_ii = ||L_i|| and
          # A_ij = L_ij/R_ii. Here ||...|| denotes Euclidean norm. Direct
          # computation shows that the jacobian determinant for the ith row is
          # R^{i-1} / A_ii.
          std_dev_sample = tf.linalg.norm(
              cholesky_sigma_sample, axis=-1
          )
          cholesky_corr_sample = (
              cholesky_sigma_sample / std_dev_sample[..., tf.newaxis]
          )
  
          cholesky_corr_std_dev_sample = (
              std_dev_sample,
              cholesky_corr_sample,
          )
          cholesky_corr_std_dev_log_prob = (
              cholesky_sigma_log_prob
              + tf.reduce_sum(
                  tf.range(dimension, dtype=dtype)
                  * tf.math.log(std_dev_sample)
                  - tf.math.log(
                      tf.linalg.diag_part(cholesky_corr_sample)
                  ),
                  axis=-1,
              )
          )
  
          return (
              cholesky_corr_std_dev_sample,
              cholesky_corr_std_dev_log_prob,
          )
  
      for dimension in (2, 3, 4, 5):
          rate = np.linspace(0.5, 2.0, 10, dtype=dtype)
          concentration = np.linspace(2.0, 5.0, 10, dtype=dtype)
  
          # We start with a distribution on SPD matrices given by the product of
          # LKJ and Exponential random variables.
          lkj_exponential_covariance_dist = (
              jds.JointDistributionSequential(
                  [
                      sample.Sample(
                          exponential.Exponential(rate=rate),
                          sample_shape=dimension,
                      ),
                      lkj.LKJ(
                          dimension=dimension,
                          concentration=concentration,
                      ),
                  ]
              )
          )
          x = self.evaluate(
              lkj_exponential_covariance_dist.sample(
                  10, seed=test_util.test_seed()
              )
          )
  
          # We transform a sample from the space of SPD matrices to the space of its
          # lower triangular Cholesky factors. We decompose it into the product of a
          # diagonal matrix and a Cholesky factor of a correlation matrix.
          (
              transformed_x,
              transformed_log_prob,
          ) = _get_transformed_sample_and_log_prob(
              x, lkj_exponential_covariance_dist.log_prob(x), dimension
          )
  
          # We now show that the transformation resulted in a distribution which
          # factors as the product of a rayleigh (the square root of an exponential)
          # and a CholeskyLKJ distribution with the same parameters as the LKJ.
          rayleigh_dist = (
              transformed_distribution.TransformedDistribution(
                  bijector=invert.Invert(square.Square()),
                  distribution=exponential.Exponential(rate=rate),
              )
          )
          cholesky_lkj_rayleigh_dist = jds.JointDistributionSequential(
              [
                  sample.Sample(rayleigh_dist, sample_shape=dimension),
                  cholesky_lkj.CholeskyLKJ(
                      dimension=dimension, concentration=concentration
                  ),
              ]
          )
          self.assertAllClose(
              self.evaluate(transformed_log_prob),
              self.evaluate(
                  cholesky_lkj_rayleigh_dist.log_prob(transformed_x)
              ),
              rtol=1e-3 if dtype == np.float32 else 1e-6,
          )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import cholesky_outer_product
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import cholesky_lkj
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lkj
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 2 for calling the CholeskyLKJTest obj's underlying member method code with example usage: obj.testLogProbMatchesTransformedDistribution(dtype) and returns: cholesky_corr_std_dev_sample,  cholesky_corr_std_dev_log_prob