code: def _apply_correction(
      theta,
      values,
      explicit_contribution,
      superdiag,
      diag,
      subdiag,
      inhomog_term_delta,
      t1,
      t2,
      dim,
      n_dims,
  ):
      """Applies correction for the given dimension."""
      rhs = (
          values
          - theta * explicit_contribution
          + theta * inhomog_term_delta * (t2 - t1)
      )
  
      # Make the given dimension the last one in the tensors, treat all the
      # other spatial dimensions as batch dimensions.
      perm = _get_permutation(values, n_dims, dim)
      if perm is not None:
          rhs = tf.transpose(rhs, perm)
          superdiag, diag, subdiag = (
              tf.transpose(c, perm) for c in (superdiag, diag, subdiag)
          )
  
      subdiag = -theta * subdiag * (t2 - t1)
      diag = 1 - theta * diag * (t2 - t1)
      superdiag = -theta * superdiag * (t2 - t1)
      result = tf.linalg.tridiagonal_solve(
          (superdiag, diag, subdiag),
          rhs,
          diagonals_format="sequence",
          partial_pivoting=False,
      )
  
      # Transpose back to how it was.
      if perm is not None:
          result = tf.transpose(result, perm)
      return result
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/steppers/douglas_adi.py
imports: import numpy
  import tensorflow.compat.v2
  from tf_quant_finance.math.pde.steppers.multidim_parabolic_equation_stepper import multidim_parabolic_equation_step
label_desc: Use image node_id 5 for calling a global function with example usage: _apply_correction(theta, values, explicit_contribution, superdiag, diag, subdiag, inhomog_term_delta, t1, t2, dim, n_dims) and returns: result