code: def _retrying_cholesky_jvp(max_iters, unroll, primals, tangents):
      """JVP for retrying_cholesky."""
      matrix, jitter = primals
      gmatrix, _ = tangents
      gmatrix = 0.5 * (gmatrix + tf.linalg.adjoint(gmatrix))
      output, shift = _retrying_cholesky_custom_gradient(
          matrix, jitter, max_iters, unroll
      )
      triangular_linop = tf.linalg.LinearOperatorLowerTriangular(output)
      middle = triangular_linop.solve(gmatrix, adjoint_arg=True)
      middle = triangular_linop.solve(middle, adjoint_arg=True)
      middle = tf.linalg.set_diag(
          middle, 0.5 * tf.linalg.diag_part(middle)
      )
      middle = tf.linalg.band_part(middle, -1, 0)
      return (
          (output, shift),
          (tf.linalg.matmul(output, middle), tf.zeros_like(shift)),
      )
  
file_path: probability/tensorflow_probability/python/experimental/distributions/marginal_fns.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 6 for calling a global function with example usage: _retrying_cholesky_jvp(max_iters, unroll, primals, tangents) and returns: output,  shift,  tf,  tf