code: def composite_scheme_step(
      first_scheme_steps, first_scheme, second_scheme
  ):
      """Composes two time marching schemes.
  
      Applies a step of parabolic PDE solver using `first_scheme` if number of
      performed steps is less than `first_scheme_steps`, and using `second_scheme`
      otherwise.
  
      Args:
        first_scheme_steps: A Python integer. Number of steps to apply
          `first_scheme` on.
        first_scheme: First time marching scheme (see `time_marching_scheme`
          argument of `parabolic_equation_step`).
        second_scheme: Second time marching scheme (see `time_marching_scheme`
          argument of `parabolic_equation_step`).
  
      Returns:
         Callable to be used in finite-difference PDE solvers (see fd_solvers.py).
      """
  
      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."""
          name = name or "composite_scheme_step"
  
          def scheme(*args, **kwargs):
              return tf.cond(
                  num_steps_performed < first_scheme_steps,
                  lambda: first_scheme(*args, **kwargs),
                  lambda: second_scheme(*args, **kwargs),
              )
  
          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
  
label_desc: Use image node_id 1 for calling a global function with example usage: composite_scheme_step(first_scheme_steps, first_scheme, second_scheme) and returns: step_fn, parabolic_equation_step, tf