code: def wrap_decay_fn(decay_fn: Callable) -> Callable:
      def init_fn(
          config: LRSchedulerConfig,
          optimizer: Optimizer,
          step_info: StepInfo,
      ) -> LambdaLR:
          return LambdaLR(
              optimizer,
              lambda current_step: decay_fn(
                  current_step,
                  step_info.num_training_steps,
                  step_info.num_warmup_steps,
                  config,
              ),
          )
  
      return init_fn
  
imports: import logging
  import math
  from typing import Any, Callable, Dict
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, LambdaLR, ReduceLROnPlateau, SequentialLR
  from ludwig.constants import MINIMIZE, TRAINING, VALIDATION
  from ludwig.modules.metric_registry import get_metric_objective
  from ludwig.schema.lr_scheduler import LRSchedulerConfig
  from ludwig.utils.metric_utils import TrainerMetric
  from ludwig.utils.trainer_utils import ProgressTracker
label_desc: Use image node_id 5 for calling a global function with example usage: wrap_decay_fn(decay_fn) and returns: init_fn, LambdaLR