code: def _weighted_scheme_implicit_part(
      vec, diag, upper, lower, theta, t1, t2
  ):
      """Implicit step of the weighted implicit-explicit scheme.
  
      Args:
        vec: A real dtype `Tensor` of shape `[num_equations, num_grid_points - 2]`.
          Represents the multiplied vector. "- 2" accounts for the boundary points,
          which the time-marching schemes do not touch.
        diag: A real dtype `Tensor` of the shape
          `[num_equations, num_grid_points - 2]`. Represents the main diagonal of
          a 3-diagonal matrix of the discretized PDE.
        upper: A real dtype `Tensor` of the shape
          `[num_equations, num_grid_points - 2]`. Represents the upper diagonal of
          a 3-diagonal matrix of the discretized PDE.
        lower:  A real dtype `Tensor` of the shape
          `[num_equations, num_grid_points - 2]`. Represents the lower diagonal of
          a 3-diagonal matrix of the discretized PDE.
        theta: A Python float between 0 and 1.
        t1: Smaller of the two times defining the step.
        t2: Greater of the two times defining the step.
  
      Returns:
        A tensor of the same shape and dtype as `vec`.
      """
      multiplier = (1 - theta) * (t1 - t2)
      diag = 1 + multiplier * diag
      upper = multiplier * upper
      lower = multiplier * lower
      return tf.linalg.tridiagonal_solve(
          [upper, diag, lower],
          vec,
          diagonals_format="sequence",
          transpose_rhs=True,
          partial_pivoting=False,
      )
  
imports: import tensorflow.compat.v2
  from tf_quant_finance.math.pde.steppers.parabolic_equation_stepper import parabolic_equation_step
label_desc: Use image node_id 4 for calling a global function with example usage: _weighted_scheme_implicit_part(vec, diag, upper, lower, theta, t1, t2) and returns: tf