code: def weighted_implicit_explicit_step(theta):
      """Creates a stepper function with weighted implicit-explicit scheme.
  
      Given a space-discretized equation
  
      ```
      du/dt = A(t) u(t) + b(t)
      ```
      (here `u` is a value vector, `A` and `b` are the matrix and the vector defined
      by the PDE), the scheme approximates the right-hand side as a weighted average
      of values taken before and after a time step:
  
      ```
      (u(t2) - u(t1)) / (t2 - t1) = theta * (A(t1) u(t1) + b(t1))
         + (1 - theta) (A(t2) u(t2) + b(t2)).
      ```
  
      Includes as particular cases the implicit (`theta = 0`), explicit
      (`theta = 1`), and Crank-Nicolson (`theta = 0.5`) schemes.
  
      The scheme is stable for `theta >= 0.5`, is second order accurate if
      `theta = 0.5` (i.e. in Crank-Nicolson case), and first order accurate
      otherwise.
  
      More details can be found in `weighted_implicit_explicit_scheme` below.
  
      Args:
        theta: A float in range `[0, 1]`. A parameter used to mix implicit and
          explicit schemes together. Value of `0.0` corresponds to the fully
          implicit scheme, `1.0` to the fully explicit, and `0.5` to the
          Crank-Nicolson scheme.
  
      Returns:
        Callable to be used in finite-difference PDE solvers (see fd_solvers.py).
      """
      scheme = weighted_implicit_explicit_scheme(theta)
  
      def step_fn(
          time,
          next_time,
          coord_grid,
          value_grid,
          boundary_conditions,
          second_order_coeff_fn,
          first_order_coeff_fn,
          zeroth_order_coeff_fn,
          inner_second_order_coeff_fn,
          inner_first_order_coeff_fn,
          num_steps_performed,
          dtype=None,
          name=None,
      ):
          """Performs the step."""
          del num_steps_performed
          name = name or "weighted_implicit_explicit_scheme"
          return parabolic_equation_step(
              time,
              next_time,
              coord_grid,
              value_grid,
              boundary_conditions,
              second_order_coeff_fn,
              first_order_coeff_fn,
              zeroth_order_coeff_fn,
              inner_second_order_coeff_fn,
              inner_first_order_coeff_fn,
              time_marching_scheme=scheme,
              dtype=dtype,
              name=name,
          )
  
      return step_fn
  
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 1 for calling a global function with example usage: weighted_implicit_explicit_step(theta) and returns: step_fn, parabolic_equation_step