code: def test_bijector_cholesky_lkj(self):
      # Let's try with a shape-shifting underlying bijector vec=>mat.
      d = sample.Sample(
          cholesky_lkj.CholeskyLKJ(4, concentration=tf.ones([5])), 2
      )
      b = d.experimental_default_event_space_bijector()
      y = self.evaluate(d.sample(seed=test_util.test_seed()))
      x = b.inverse(y) + 0
      self.assertAllClose(y, b.forward(x))
      y = self.evaluate(d.sample(7, seed=test_util.test_seed()))
      x = b.inverse(y) + 0
      self.assertAllClose(y, b.forward(x))
      d2 = independent.Independent(
          cholesky_lkj.CholeskyLKJ(4, concentration=tf.ones([5, 2])),
          reinterpreted_batch_ndims=1,
      )
      b2 = d2.experimental_default_event_space_bijector()
      self.assertAllClose(
          b2.forward_log_det_jacobian(x, event_ndims=len([2, 6])),
          b.forward_log_det_jacobian(x, event_ndims=len([2, 6])),
      )
      x_sliced = x[..., :1, :]
      x_bcast = tf.concat([x_sliced, x_sliced], axis=-2)
      self.assertAllClose(
          b2.forward_log_det_jacobian(x_bcast, event_ndims=len([2, 6])),
          b.forward_log_det_jacobian(x_sliced, event_ndims=len([1, 6])),
      )
      # Should this test pass? Right now, Independent's default bijector does not
      # broadcast the LDJ to match underlying batch shape.
      # self.assertAllClose(
      #     b2.forward_log_det_jacobian(x_sliced, event_ndims=len([1, 6])),
      #     b.forward_log_det_jacobian(x_bcast, event_ndims=len([2, 6])))
      self.assertAllClose(
          b2.inverse_log_det_jacobian(y, event_ndims=len([2, 4, 4])),
          b.inverse_log_det_jacobian(y, event_ndims=len([2, 4, 4])),
      )
      y_sliced = y[..., :1, :, :]
      y_bcast = tf.concat([y_sliced, y_sliced], axis=-3)
      self.assertAllClose(
          b2.inverse_log_det_jacobian(
              y_bcast, event_ndims=len([2, 4, 4])
          ),
          b.inverse_log_det_jacobian(
              y_sliced, event_ndims=len([1, 4, 4])
          ),
      )
      # Should this test pass? Right now, Independent's default bijector does not
      # broadcast the LDJ to match underlying batch shape.
      # self.assertAllClose(
      #     b2.inverse_log_det_jacobian(y_sliced, event_ndims=len([1, 4, 4])),
      #     b.inverse_log_det_jacobian(y_bcast, event_ndims=len([2, 4, 4])))
      self.assertAllClose(
          b.forward_log_det_jacobian(x_sliced, event_ndims=len([2, 6])),
          -b.inverse_log_det_jacobian(
              y_bcast, event_ndims=len([2, 4, 4])
          ),
          rtol=1e-5,
      )
  
      # Now, with another sample shape.
      d = sample.Sample(
          cholesky_lkj.CholeskyLKJ(4, concentration=tf.ones([5])),
          [2, 7],
      )
      b = d.experimental_default_event_space_bijector()
      y = self.evaluate(d.sample(11, seed=test_util.test_seed()))
      x = b.inverse(y) + 0
      self.assertAllClose(y, b.forward(x))
      d2 = independent.Independent(
          cholesky_lkj.CholeskyLKJ(4, concentration=tf.ones([5, 2, 7])),
          reinterpreted_batch_ndims=2,
      )
      b2 = d2.experimental_default_event_space_bijector()
      self.assertAllClose(
          b2.forward_log_det_jacobian(x, event_ndims=len([2, 7, 6])),
          b.forward_log_det_jacobian(x, event_ndims=len([2, 7, 6])),
      )
      self.assertAllClose(
          b2.inverse_log_det_jacobian(y, event_ndims=len([2, 7, 4, 4])),
          b.inverse_log_det_jacobian(y, event_ndims=len([2, 7, 4, 4])),
      )
  
      # Now, with another batch shape.
      d = sample.Sample(
          cholesky_lkj.CholeskyLKJ(4, concentration=tf.ones([5, 7])), 2
      )
      b = d.experimental_default_event_space_bijector()
      y = self.evaluate(d.sample(11, seed=test_util.test_seed()))
      x = b.inverse(y) + 0
      self.assertAllClose(y, b.forward(x))
      self.assertAllClose(
          b2.forward_log_det_jacobian(x, event_ndims=len([2, 7, 6])),
          b.forward_log_det_jacobian(x, event_ndims=len([2, 7, 6])),
      )
      self.assertAllClose(
          b2.inverse_log_det_jacobian(y, event_ndims=len([2, 7, 4, 4])),
          b.inverse_log_det_jacobian(y, event_ndims=len([2, 7, 4, 4])),
      )
  
      # Also verify it properly handles an extra "event" dim.
      self.assertAllClose(
          b2.forward_log_det_jacobian(x, event_ndims=len([5, 2, 7, 6])),
          b.forward_log_det_jacobian(x, event_ndims=len([5, 2, 7, 6])),
      )
      self.assertAllClose(
          b2.inverse_log_det_jacobian(
              y, event_ndims=len([5, 2, 7, 4, 4])
          ),
          b.inverse_log_det_jacobian(
              y, event_ndims=len([5, 2, 7, 4, 4])
          ),
      )
  
imports: import functools
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import scale_matvec_tril
  from tensorflow_probability.python.distributions import cholesky_lkj
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.distributions import logistic
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 17 for calling the SampleDistributionTest obj's underlying member method code with example usage: obj.test_bijector_cholesky_lkj() without return types