code: def make_cholesky_with_jitter_fn(jitter=1e-6):
      """Make a function that adds diagonal jitter before Cholesky factoring.
  
      Suitable for use in the `cholesky_fn` parameter for
      `GaussianProcessRegressionModelWithCholesky`.
  
      Args:
        jitter: Float diagonal jitter to add. Default value: 1e-6.
  
      Returns:
        cholesky_with_jitter: Function that computes jittered Cholesky.
      """
  
      def cholesky_with_jitter(matrix):
          jittered = tf.linalg.set_diag(
              matrix, tf.linalg.diag_part(matrix) + jitter
          )
          return tf.linalg.cholesky(jittered)
  
      return cholesky_with_jitter
  
label_desc: Use image node_id 1 for calling a global function with example usage: make_cholesky_with_jitter_fn(jitter) and returns: cholesky_with_jitter, tf