code: def run_train_steps(
      m: "ModelWithCustomSteps",
      train_steps: List["CustomTrainStep"],
      *,
      batch_idx: int,
      batch: tensor_dict_type,
      trainer: "ITrainer",
      forward_kwargs: Dict[str, Any],
      loss_kwargs: Dict[str, Any],
  ) -> StepOutputs:
      """
      Runs a series of custom training steps on a given model and batch of data.
  
      Parameters
      ----------
      m : ModelWithCustomSteps
          The model to train.
      train_steps : List[CustomTrainStep]
          The custom training steps to run.
      batch_idx : int
          The current batch index.
      batch : tensor_dict_type
          The batch of data to use for training.
      trainer : ITrainer
          The trainer object used to train the model.
      forward_kwargs : Dict[str, Any]
          Additional arguments to pass to the forward pass of the model.
      loss_kwargs : Dict[str, Any]
          Additional arguments to pass to the loss function of each training step.
  
      Returns
      -------
      StepOutputs
          An object containing the outputs of the forward pass and the calculated loss values of the training steps.
  
      Step by step explanation
      ------------------------
      1. Initialize variables: `forward` (an empty dictionary), `loss_dict` (an empty dictionary), `any_update`
      (a bool flag set to `False`), `performed_scheduler_step` (a bool flag set to `False`), and `update_fn` (a
      function returned by the `get_update_fn` function defined above).
      2. Check whether the forward pass should have gradients (`fw_has_grad`) and which training step to use for the
      forward pass (`fw_train_step`). This is done by looping through each training step and checking its
      `requires_new_forward` and `requires_grad_in_forward` attributes.
      3. If `fw_has_grad` is `False` and a subsequent training step requires gradients in the forward pass, raise a
      ValueError with a message indicating which training steps have conflicting requirements.
      4. Loop through each training step and execute the following steps for each:
        1) Check whether the current training step should be skipped. If so, move on to the next training step.
        2) If this is the first training step, or if `requires_new_forward` is `True` for the current training step,
        execute the forward pass of the model and store the output in `forward`. The `no_grad_context` context manager
        is used to prevent gradients from being calculated if `requires_grad_in_forward` is `False`.
        3) Get the optimizer to be used for this training step.
        4) If `enable_toggle_optimizer` is `True` for this training step, temporarily switch to the optimizer associated
        with this training step using the `toggle_optimizer` context manager.
        5) Calculate the loss for this training step using the model, state, batch, and forward pass outputs. The
        `autocast` context manager is used if mixed-precision training is enabled.
        6) Update the optimizer if `train_step.grad_accumulate` is a factor of the current `state.step`.
        7) Update the `loss_dict` with the loss values for this training step.
        8) If an optimizer update occurred, set `any_update` to `True`, and if `requires_scheduler_step` is `True` for
        this training step, call `trainer.scheduler_step()` to update the learning rate.
      5. If any optimizer updates occurred but no scheduler steps were performed, call `trainer.scheduler_step()` to
      update the learning rate.
      6. Loop through each training step and call its callback function with the model, trainer, batch, and forward pass outputs.
      7. Return the `StepOutputs` object containing the forward pass outputs and loss values.
      """
  
      state = trainer.state
      forward: Union[tensor_dict_type, List[tensor_dict_type]] = {}
      loss_dict = {}
      update_fn = get_update_fn(trainer)
      any_update = False
      performed_scheduler_step = False
      # sanity check
      fw_has_grad = True
      fw_train_step: Any = ()
      for i, train_step in enumerate(train_steps):
          if i == 0 or train_step.requires_new_forward:
              fw_has_grad = train_step.requires_grad_in_forward
              fw_train_step = train_step
          if not fw_has_grad and train_step.requires_grad_in_forward:
              fw_name = fw_train_step.__class__.__name__
              current_name = train_step.__class__.__name__
              raise ValueError(
                  f"current forward pass comes from '{fw_name}' and has no grad, "
                  f"but '{current_name}' requires grad in forward. You can either set "
                  f"`requires_grad_in_forward` of '{fw_name}' to True, or set "
                  f"`requires_new_forward` of '{current_name}' to True."
              )
      # run train steps
      get_fw = lambda: m.run(batch_idx, batch, state, **forward_kwargs)
      for i, train_step in enumerate(train_steps):
          if train_step.should_skip(m, state):
              continue
          if i == 0 or train_step.requires_new_forward:
              with no_grad_context(
                  enabled=not train_step.requires_grad_in_forward
              ):
                  if train_step.num_forward == 1:
                      forward = get_fw()
                  else:
                      forward = [
                          get_fw()
                          for _ in range(train_step.num_forward)
                      ]
          optimizer = trainer.optimizers[train_step.scope]
          with toggle_optimizer(
              m, optimizer, enabled=train_step.enable_toggle_optimizer
          ):
              with autocast(
                  enabled=trainer.config.mixed_precision
                  != PrecisionType.NO
              ):
                  loss_res = train_step.loss_fn(
                      m, state, batch, forward, **loss_kwargs
                  )
              update = state.step % train_step.grad_accumulate == 0
              update_fn(loss_res.loss, optimizer, update)
              loss_dict.update(loss_res.losses)
          if update:
              any_update = True
              performed_scheduler_step = (
                  train_step.requires_scheduler_step
              )
              if performed_scheduler_step:
                  trainer.scheduler_step()
      if any_update and not performed_scheduler_step:
          trainer.scheduler_step()
      # callbacks
      for train_step in train_steps:
          train_step.callback(m, trainer, batch, forward)
      return StepOutputs(forward, loss_dict)
  
imports: import torch.nn
  from abc import abstractmethod
  from abc import ABC
  from abc import ABCMeta
  from torch import Tensor
  from typing import Any
  from typing import Dict
  from typing import List
  from typing import Union
  from typing import Callable
  from typing import Optional
  from typing import NamedTuple
  from torch.optim import Optimizer
  from cftool.types import tensor_dict_type
  from torch.cuda.amp import autocast
  from ...schema import IDLModel
  from ...schema import ITrainer
  from ...schema import IDataLoader
  from ...schema import StepOutputs
  from ...schema import TrainerState
  from ...schema import PrecisionType
  from ...schema import TrainerConfig
  from ...schema import MetricsOutputs
  from ...constants import LOSS_KEY
  from ...data.utils import TensorBatcher
  from ...misc.toolkit import get_device
  from ...misc.toolkit import no_grad_context
  from ...misc.toolkit import toggle_optimizer
label_desc: Use image node_id 2 for calling a global function with example usage: run_train_steps(m, train_steps) and returns: StepOutputs