code: def add_diagonal_shift(matrix, shift):
      broadcast_shape = distribution_util.get_broadcast_shape(
          matrix, shift[..., tf.newaxis]
      )
      matrix = tf.broadcast_to(matrix, broadcast_shape)
      return tf.linalg.set_diag(
          matrix,
          tf.linalg.diag_part(matrix) + shift,
          name="add_diagonal_shift",
      )
  
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 4 for calling a global function with example usage: add_diagonal_shift(matrix, shift) and returns: tf