code: def compute_kernel_matrix(
      kernel, index_points, observation_noise_variance
  ):
      kernel_matrix = kernel.matrix(index_points, index_points)
      observation_noise_variance = tf.convert_to_tensor(
          observation_noise_variance
      )
      # We are compute K + obs_noise_variance * I. The shape of this matrix is
      # going to be a broadcast of the shapes of K and obs_noise_variance * I.
      broadcast_shape = distribution_util.get_broadcast_shape(
          kernel_matrix,
          # We pad with two single dimension since this represents a batch of
          # scaled identity matrices.
          observation_noise_variance[..., tf.newaxis, tf.newaxis],
      )
      kernel_matrix = tf.broadcast_to(kernel_matrix, broadcast_shape)
      return add_diagonal_shift(
          kernel_matrix, observation_noise_variance[..., tf.newaxis]
      )
  
file_path: probability/tensorflow_probability/python/distributions/internal/stochastic_process_util.py
imports: import functools
  import tensorflow.compat.v2
  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.psd_kernels.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 12 for calling a global function with example usage: compute_kernel_matrix(kernel, index_points, observation_noise_variance) and returns: add_diagonal_shift