code: def _weighted_scheme_explicit_part(
      vec, diag, upper, lower, theta, t1, t2
  ):
      """Explicit 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 = theta * (t2 - t1)
      diag = 1 + multiplier * diag
      upper = multiplier * upper
      lower = multiplier * lower
  
      # Multiply the tridiagonal matrix by the vector.
      diag_part = diag * vec
      zeros = tf.zeros_like(lower[..., :1])
      lower_part = tf.concat(
          (zeros, lower[..., 1:] * vec[..., :-1]), axis=-1
      )
      upper_part = tf.concat(
          (upper[..., :-1] * vec[..., 1:], zeros), axis=-1
      )
      return lower_part + diag_part + upper_part
  
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 3 for calling a global function with example usage: _weighted_scheme_explicit_part(vec, diag, upper, lower, theta, t1, t2) and returns: unknown