code: def get_param_info(optim: Optimizer):
      # Get a backup of necessary information of parameters for future use, which includes:
      # 1. A mapping from integer param_id to param32 shape.
  
      if optim is None:
          return {}
      param_info = {"id2shape": {}}
      start_index = 0
      for group in optim.param_groups:
          for param_id, param in enumerate(
              group["params"], start_index
          ):
              original_shape = (
                  param.shape
                  if isinstance(param, torch.Tensor)
                  else None
              )
              param_info["id2shape"][param_id] = original_shape
  
          start_index += len(group["params"])
  
      return param_info
  
imports: import gc
  import logging
  import os
  import random
  from pathlib import Path
  from typing import Callable, Iterator, List, Optional, Tuple
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO
  from colossalai.checkpoint_io.utils import get_model_base_filenames, get_optimizer_base_filenames, load_shard_state_dict, save_config_file, save_state_dict, save_state_dict_shards
  from colossalai.cluster import DistCoordinator, ProcessGroupMesh
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
  from colossalai.zero import GeminiDDP, GeminiOptimizer
  from colossalai.zero.gemini.memory_tracer import MemStats
  from .dp_plugin_base import DPPluginBase
label_desc: Use image node_id 1 for calling a global function with example usage: get_param_info(optim) and returns: param_info, dict