code: def _compute_flattened_scale(
      kernel, index_points, cholesky_fn, observation_noise_variance=None
  ):
      """Computes a matrix square root of the flattened covariance matrix.
  
      Given a multi-task kernel `k`, computes a matrix square root of the
      matrix over all tasks of `index_points`. That is, compute `S` such that
      `S^T @ S = k.matrix_over_all_tasks(index_points, index_points)`.
  
      In the case of a `Separable` or `Independent` kernel, this function tries to
      do this efficiently in O(N^3 + T^3) time where `N` is the number of
      `index_points` and `T` is the number of tasks.
  
      Args:
        kernel: `MultiTaskKernel`-like instance representing the GP's covariance
          function.
        index_points: (Nested) `float` `Tensor` representing finite collection, or
          batch of collections, of points in the index set over which the GP is
          defined. Shape (or shape of each nested component) has the form
          `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
          dimensions and must equal `kernel.feature_ndims` (or its corresponding
          nested component) and `e` is the number (size) of index points in each
          batch. Ultimately this distribution corresponds to an `e`-dimensional
          multivariate normal. The batch shape must be broadcastable with
          `kernel.batch_shape`.
        cholesky_fn: Callable which takes a single (batch) matrix argument and
          returns a Cholesky-like lower triangular factor.  Default value: `None`,
          in which case `make_cholesky_with_jitter_fn(1e-6)` is used.
        observation_noise_variance: `float` `Tensor` representing the variance
          of the noise in the Normal likelihood distribution of the model. May be
          batched, in which case the batch shape must be broadcastable with the
          shapes of all other batched parameters (`kernel.batch_shape`,
          `index_points`, etc.).
          Default value: `None`
      Returns:
        scale_operator: `LinearOperator` representing a matrix square root of
        the flattened kernel matrix over all tasks.
  
      """
      # This is of shape KN x KN, where K is the number of outputs
      kernel_matrix = kernel.matrix_over_all_tasks(
          index_points, index_points
      )
      if observation_noise_variance is None:
          return cholesky_util.cholesky_from_fn(
              kernel_matrix, cholesky_fn
          )
  
      observation_noise_variance = tf.convert_to_tensor(
          observation_noise_variance
      )
  
      # We can add the observation noise to each block.
      if isinstance(kernel, multitask_kernel.Independent):
          # The Independent kernel matrix is realized as a kronecker product of the
          # kernel over inputs, and an identity matrix per task (representing
          # independent tasks). Update the diagonal of the first matrix and take the
          # cholesky of it (since the cholesky of the second matrix will remain the
          # identity matrix.)
          base_kernel_matrix = kernel_matrix.operators[0].to_dense()
  
          broadcast_shape = distribution_util.get_broadcast_shape(
              base_kernel_matrix,
              observation_noise_variance[..., tf.newaxis, tf.newaxis],
          )
          base_kernel_matrix = tf.broadcast_to(
              base_kernel_matrix, broadcast_shape
          )
          base_kernel_matrix = tf.linalg.set_diag(
              base_kernel_matrix,
              tf.linalg.diag_part(base_kernel_matrix)
              + observation_noise_variance[..., tf.newaxis],
          )
          base_kernel_matrix = tf.linalg.LinearOperatorFullMatrix(
              base_kernel_matrix
          )
          kernel_matrix = tf.linalg.LinearOperatorKronecker(
              operators=[base_kernel_matrix]
              + kernel_matrix.operators[1:]
          )
          return cholesky_util.cholesky_from_fn(
              kernel_matrix, cholesky_fn
          )
  
      if isinstance(kernel, multitask_kernel.Separable):
          # When `kernel_matrix` is a kronecker product, we can compute
          # an eigenvalue decomposition to get a matrix square-root, which will
          # be faster than densifying the kronecker product.
  
          # Let K = A X B. Let A (and B) have an eigenvalue decomposition of
          # U @ D @ U^T, where U is an orthogonal matrix. Then,
          # K = (U_A @ D_A @ U_A^T) X (U_B @ D_B @ U_B^T) =
          # (U_A X U_B) @ (D_A X D_B) @ (U_A X U_B)^T
          # Thus, a matrix square root of K would be
          # (U_A X U_B) @ (sqrt(D_A) X sqrt(D_B)) which offers
          # efficient matmul and solves.
  
          # Now, if we update the diagonal by `v * I`, we have
          # (U_A X U_B) @ (sqrt((D_A X D_B + vI)) @ (U_A X U_B)^T
          # which still admits an efficient matmul and solve.
  
          kronecker_diags = []
          kronecker_orths = []
          for block in kernel_matrix.operators:
              # No need to take an eigenvalue decomposition for diagonal operators since
              # they are already in factored form.
              if isinstance(
                  block,
                  (
                      tf.linalg.LinearOperatorDiag,
                      tf.linalg.LinearOperatorIdentity,
                      tf.linalg.LinearOperatorScaledIdentity,
                  ),
              ):
                  kronecker_diags.append(block)
                  kronecker_orths.append(
                      tf.linalg.LinearOperatorIdentity(
                          block.domain_dimension_tensor(),
                          dtype=block.dtype,
                      )
                  )
              else:
                  diag, orth = tf.linalg.eigh(block.to_dense())
                  kronecker_diags.append(
                      tf.linalg.LinearOperatorDiag(diag)
                  )
                  kronecker_orths.append(
                      linear_operator_unitary.LinearOperatorUnitary(
                          orth
                      )
                  )
  
          full_diag = tf.linalg.LinearOperatorKronecker(
              kronecker_diags
          ).diag_part()
          full_diag = (
              full_diag + observation_noise_variance[..., tf.newaxis]
          )
          scale_diag = tf.math.sqrt(full_diag)
          diag_operator = tf.linalg.LinearOperatorDiag(
              scale_diag,
              is_square=True,
              is_non_singular=True,
              is_positive_definite=True,
          )
  
          orthogonal_operator = tf.linalg.LinearOperatorKronecker(
              kronecker_orths, is_square=True, is_non_singular=True
          )
          # This is efficient as a scale matrix. When used for matmuls, we take
          # advantage of the kronecker product and diagonal operator. When used for
          # solves, we take advantage of the orthogonal and diagonal structure,
          # which essentially reduces to the matmul case.
          return orthogonal_operator.matmul(diag_operator)
  
      # By default densify the kernel matrix and add noise.
  
      kernel_matrix = kernel_matrix.to_dense()
      broadcast_shape = distribution_util.get_broadcast_shape(
          kernel_matrix,
          observation_noise_variance[..., tf.newaxis, tf.newaxis],
      )
      kernel_matrix = tf.broadcast_to(kernel_matrix, broadcast_shape)
      kernel_matrix = tf.linalg.set_diag(
          kernel_matrix,
          tf.linalg.diag_part(kernel_matrix)
          + observation_noise_variance[..., tf.newaxis],
      )
      kernel_matrix = tf.linalg.LinearOperatorFullMatrix(kernel_matrix)
      kernel_cholesky = cholesky_util.cholesky_from_fn(
          kernel_matrix, cholesky_fn
      )
      return kernel_cholesky
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import cholesky_util
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions.internal import stochastic_process_util
  from tensorflow_probability.python.experimental.linalg import linear_operator_unitary
  from tensorflow_probability.python.experimental.psd_kernels import multitask_kernel
  from tensorflow_probability.python.internal import batch_shape_lib
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_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 tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: _compute_flattened_scale(kernel, index_points, cholesky_fn, observation_noise_variance) and returns: kernel_cholesky, cholesky_util, cholesky_util, orthogonal_operator