code: def _compute_observation_scale(
      kernel,
      observation_index_points,
      cholesky_fn,
      observation_noise_variance=None,
      observations_is_missing=None,
  ):
      """Compute matrix square root of the kernel on observation index points."""
      if observations_is_missing is not None:
          observations_is_missing = tf.convert_to_tensor(
              observations_is_missing
          )
          # If observations are missing, there's nothing we can do to preserve the
          # operator structure, so densify.
  
          observation_covariance = kernel.matrix_over_all_tasks(
              observation_index_points, observation_index_points
          ).to_dense()
  
          if observation_noise_variance is not None:
              broadcast_shape = distribution_util.get_broadcast_shape(
                  observation_covariance,
                  observation_noise_variance[
                      ..., tf.newaxis, tf.newaxis
                  ],
              )
              observation_covariance = tf.broadcast_to(
                  observation_covariance, broadcast_shape
              )
              observation_covariance = _add_diagonal_shift(
                  observation_covariance, observation_noise_variance
              )
          vec_observations_is_missing = _vec(observations_is_missing)
          observation_covariance = tf.linalg.LinearOperatorFullMatrix(
              psd_kernels_util.mask_matrix(
                  observation_covariance,
                  is_missing=vec_observations_is_missing,
              ),
              is_non_singular=True,
              is_positive_definite=True,
          )
          observation_scale = cholesky_util.cholesky_from_fn(
              observation_covariance, cholesky_fn
          )
      else:
          observation_scale = mtgp._compute_flattened_scale(  # pylint:disable=protected-access
              kernel=kernel,
              index_points=observation_index_points,
              cholesky_fn=cholesky_fn,
              observation_noise_variance=observation_noise_variance,
          )
  
      return observation_scale
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  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.distributions import multitask_gaussian_process
  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_probability.python.math.psd_kernels.internal import util
label_desc: Use image node_id 5 for calling a global function with example usage: _compute_observation_scale(kernel, observation_index_points, cholesky_fn, observation_noise_variance, observations_is_missing) and returns: observation_scale