code: def _lr_stages(stages: List[WarmupStage]) -> List[WarmupStage]:
      last_stage = WarmupStage(
          policy=WarmupPolicy.NONE, max_iters=1 << 63, value=1.0
      )
      if len(stages) == 0:
          return [last_stage]
  
      start_iter = 0
      for stage in stages:
          assert stage.max_iters > start_iter, (
              f"Max iter of the stage {stage} must be greater than the previous "
              f"max iter {start_iter}"
          )
          start_iter = stage.max_iters
          if stage.decay_iters <= 0:
              if stage.policy == WarmupPolicy.STEP:
                  stage.decay_iters = 1
              else:
                  stage.decay_iters = stage.max_iters
      return stages + [last_stage]
  
file_path: torchrec/torchrec/optim/warmup.py
imports: import logging
  import math
  from dataclasses import dataclass
  from enum import Enum, unique
  from typing import Any, List, Tuple
  import torch
  from torchrec.optim.keyed import KeyedOptimizer, OptimizerWrapper
label_desc: Use image node_id 1 for calling a global function with example usage: _lr_stages(stages) and returns: unknown, list