code: def sample_lkj(
      num_samples,
      dimension,
      concentration,
      cholesky_space=False,
      seed=None,
      name=None,
  ):
      """Returns a Tensor of samples from an LKJ distribution.
  
      Args:
        num_samples: Python `int`. The number of samples to draw.
        dimension: Python `int`. The dimension of correlation matrices.
        concentration: `Tensor` representing the concentration of the LKJ
          distribution.
        cholesky_space: Python `bool`. Whether to take samples from LKJ or
          Chol(LKJ).
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        name: Python `str` name prefixed to Ops created by this function.
  
      Returns:
        samples: A Tensor of correlation matrices (or Cholesky factors of
          correlation matrices if `cholesky_space = True`) with shape
          `[n] + B + [D, D]`, where `B` is the shape of the `concentration`
          parameter, and `D` is the `dimension`.
  
      Raises:
        ValueError: If `dimension` is negative.
      """
      if dimension < 0:
          raise ValueError(
              "Cannot sample negative-dimension correlation matrices."
          )
      # Notation below: B is the batch shape, i.e., tf.shape(concentration)
  
      with tf.name_scope("sample_lkj" or name):
          concentration = tf.convert_to_tensor(concentration)
          if not dtype_util.is_floating(concentration.dtype):
              raise TypeError(
                  "The concentration argument should have floating type, not "
                  "{}".format(dtype_util.name(concentration.dtype))
              )
  
          batch_shape = ps.concat(
              [[num_samples], ps.shape(concentration)], axis=0
          )
          dtype = concentration.dtype
          if dimension <= 1:
              # For any dimension <= 1, there is only one possible correlation matrix.
              shape = ps.concat(
                  [batch_shape, [dimension, dimension]], axis=0
              )
              return tf.ones(shape=shape, dtype=dtype)
  
          # We need 1 seed for beta and 1 seed for tril_spherical_uniform.
          beta_seed, tril_spherical_uniform_seed = samplers.split_seed(
              seed, n=2, salt="sample_lkj"
          )
  
          # Note that the sampler below deviates from [1], by doing the sampling in
          # cholesky space. This does not change the fundamental logic of the
          # sampler, but does speed up the sampling.
          # In addition, we also vectorize the computation to make the sampler
          # more feasible to use in problems where `dimension` is large.
  
          beta_conc = concentration + (dimension - 2.0) / 2.0
          dimension_range = np.arange(
              1.0, dimension, dtype=dtype_util.as_numpy_dtype(dtype)
          )
          beta_conc1 = dimension_range / 2.0
          beta_conc0 = (
              beta_conc[..., tf.newaxis] - (dimension_range - 1) / 2.0
          )
          beta_dist = beta.Beta(
              concentration1=beta_conc1, concentration0=beta_conc0
          )
          # norm is y in reference [1].
          norm = beta_dist.sample(
              sample_shape=[num_samples], seed=beta_seed
          )
          # distance shape: B + [dimension - 1, 1] for broadcast
          distance = tf.sqrt(norm)[..., tf.newaxis]
  
          # direction is u in reference [1].
          # direction follows the spherical uniform distribution and will be stored
          # in a lower triangular matrix, hence it will have shape:
          # B + [dimension - 1, dimension - 1]
          direction = _tril_spherical_uniform(
              dimension - 1,
              batch_shape,
              dtype,
              tril_spherical_uniform_seed,
          )
  
          # raw_correlation is w in reference [1].
          # shape: B + [dimension - 1, dimension - 1]
          raw_correlation = distance * direction
  
          # This is the rows in the cholesky of the result,
          # which differs from the construction in reference [1].
          # In the reference, the new row `z` = chol_result @ raw_correlation^T
          # = C @ raw_correlation^T (where as short hand we use C = chol_result).
          # We prove that the below equation is the right row to add to the
          # cholesky, by showing equality with reference [1].
          # Let S be the sample constructed so far, and let `z` be as in
          # reference [1]. Then at this iteration, the new sample S' will be
          # [[S z^T]
          #  [z 1]]
          # In our case we have the cholesky decomposition factor C, so
          # we want our new row x (same size as z) to satisfy:
          #  [[S z^T]  [[C 0]    [[C^T  x^T]         [[CC^T  Cx^T]
          #   [z 1]] =  [x k]]    [0     k]]  =       [xC^t   xx^T + k**2]]
          # Since C @ raw_correlation^T = z = C @ x^T, and C is invertible,
          # we have that x = raw_correlation. Also 1 = xx^T + k**2, so k
          # = sqrt(1 - xx^T) = sqrt(1 - |raw_correlation|**2) = sqrt(1 -
          # distance**2).
          paddings_prepend = [[0, 0]] * len(batch_shape)
          diag = tf.pad(
              tf.sqrt(1.0 - norm),
              paddings_prepend + [[1, 0]],
              constant_values=1.0,
          )
          chol_result = tf.pad(
              raw_correlation,
              paddings_prepend + [[1, 0], [0, 1]],
              constant_values=0.0,
          )
          chol_result = tf.linalg.set_diag(chol_result, diag)
  
          if cholesky_space:
              return chol_result
  
          result = tf.matmul(chol_result, chol_result, transpose_b=True)
          # The diagonal for a correlation matrix should always be ones. Due to
          # numerical instability the matmul might not achieve that, so manually set
          # these to ones.
          result = tf.linalg.set_diag(
              result,
              tf.ones(shape=ps.shape(result)[:-1], dtype=result.dtype),
          )
          # This sampling algorithm can produce near-PSD matrices on which standard
          # algorithms such as `tf.linalg.cholesky` or
          # `tf.linalg.self_adjoint_eigvals` fail. Specifically, as documented in
          # b/116828694, around 2% of trials of 900,000 5x5 matrices (distributed
          # according to 9 different concentration parameter values) contained at
          # least one matrix on which the Cholesky decomposition failed.
          return result
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import cholesky_outer_product
  from tensorflow_probability.python.bijectors import correlation_cholesky
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import linalg
  from tensorflow_probability.python.math import special
  from tensorflow_probability.python.math.numeric import clip_by_value_preserve_gradient
  from tensorflow_probability.python.random import random_ops
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 2 for calling a global function with example usage: sample_lkj(num_samples, dimension, concentration, cholesky_space, seed, name) and returns: result, tf, chol_result