code: def schedule(
      *,
      wait: int,
      warmup: int,
      active: int,
      repeat: int = 0,
      skip_first: int = 0,
  ) -> Callable:
      """
      Returns a callable that can be used as profiler ``schedule`` argument. The profiler will skip
      the first ``skip_first`` steps, then wait for ``wait`` steps, then do the warmup for the next ``warmup`` steps,
      then do the active recording for the next ``active`` steps and then repeat the cycle starting with ``wait`` steps.
      The optional number of cycles is specified with the ``repeat`` parameter, the zero value means that
      the cycles will continue until the profiling is finished.
      """
  
      def schedule_fn(step: int) -> ProfilerAction:
          assert step >= 0
          if step < skip_first:
              return ProfilerAction.NONE
          else:
              step -= skip_first
          num_steps = wait + warmup + active
          if repeat > 0 and step / num_steps >= repeat:
              return ProfilerAction.NONE
          mod_step = step % num_steps
          if mod_step < wait:
              return ProfilerAction.NONE
          elif mod_step < wait + warmup:
              return ProfilerAction.WARMUP
          else:
              return (
                  ProfilerAction.RECORD
                  if mod_step < num_steps - 1
                  else ProfilerAction.RECORD_AND_SAVE
              )
  
      assert (
          wait >= 0
          and warmup >= 0
          and active > 0
          and repeat >= 0
          and skip_first >= 0
      ), "Invalid profiler schedule arguments"
      if warmup == 0:
          warn(
              "Profiler won't be using warmup, this can skew profiler results"
          )
      return schedule_fn
  
imports: import gzip
  import json
  import os
  import tempfile
  from enum import Enum
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
  from warnings import warn
  import torch
  import torch.autograd.profiler
  from torch._C import _get_privateuse1_backend_name
  from torch._C._profiler import _add_execution_trace_observer, _disable_execution_trace_observer, _enable_execution_trace_observer, _ExperimentalConfig, _remove_execution_trace_observer
  from torch.autograd import kineto_available, ProfilerActivity
  from torch.profiler._memory_profiler import MemoryProfile, MemoryProfileTimeline
label_desc: Use image node_id 2 for calling a global function with example usage: schedule() and returns: schedule_fn, ProfilerAction, ProfilerAction, ProfilerAction, ProfilerAction, unknown