code: def cholesky_concat(chol, cols, name=None):
      """Concatenates `chol @ chol.T` with additional rows and columns.
  
      This operation is conceptually identical to:
      ```python
      def cholesky_concat_slow(chol, cols):  # cols shaped (n + m) x m = z x m
        mat = tf.matmul(chol, chol, adjoint_b=True)  # shape of n x n
        # Concat columns.
        mat = tf.concat([mat, cols[..., :tf.shape(mat)[-2], :]], axis=-1)  # n x z
        # Concat rows.
        mat = tf.concat([mat, tf.linalg.matrix_transpose(cols)], axis=-2)  # z x z
        return tf.linalg.cholesky(mat)
      ```
      but whereas `cholesky_concat_slow` would cost `O(z**3)` work,
      `cholesky_concat` only costs `O(z**2 + m**3)` work.
  
      The resulting (implicit) matrix must be symmetric and positive definite.
      Thus, the bottom right `m x m` must be self-adjoint, and we do not require a
      separate `rows` argument (which can be inferred from `conj(cols.T)`).
  
      Args:
        chol: Cholesky decomposition of `mat = chol @ chol.T`.
        cols: The new columns whose first `n` rows we would like concatenated to the
          right of `mat = chol @ chol.T`, and whose conjugate transpose we would
          like concatenated to the bottom of `concat(mat, cols[:n,:])`. A `Tensor`
          with final dims `(n+m, m)`. The first `n` rows are the top right rectangle
          (their conjugate transpose forms the bottom left), and the bottom `m x m`
          is self-adjoint.
        name: Optional name for this op.
  
      Returns:
        chol_concat: The Cholesky decomposition of:
          ```
          [ [ mat  cols[:n, :] ]
            [   conj(cols.T)   ] ]
          ```
      """
      with tf.name_scope(name or "cholesky_extend"):
          dtype = dtype_util.common_dtype(
              [chol, cols], dtype_hint=tf.float32
          )
          chol = tf.convert_to_tensor(chol, name="chol", dtype=dtype)
          cols = tf.convert_to_tensor(cols, name="cols", dtype=dtype)
          n = ps.shape(chol)[-1]
          mat_nm, mat_mm = cols[..., :n, :], cols[..., n:, :]
          solved_nm = tf.linalg.triangular_solve(chol, mat_nm)
          lower_right_mm = tf.linalg.cholesky(
              mat_mm - tf.matmul(solved_nm, solved_nm, adjoint_a=True)
          )
          lower_left_mn = tf.math.conj(
              tf.linalg.matrix_transpose(solved_nm)
          )
          out_batch = ps.shape(solved_nm)[:-2]
          chol = tf.broadcast_to(
              chol, ps.concat([out_batch, ps.shape(chol)[-2:]], axis=0)
          )
          top_right_zeros_nm = tf.zeros_like(solved_nm)
          return tf.concat(
              [
                  tf.concat([chol, top_right_zeros_nm], axis=-1),
                  tf.concat([lower_left_mn, lower_right_mm], axis=-1),
              ],
              axis=-2,
          )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 1 for calling a global function with example usage: cholesky_concat(chol, cols, name) and returns: tf