code: def safe_get_full_grad(param):
      """Assemble and return the fp32 gradient of a low-precision (e.g., fp16) parameter.
  
      Args:
          param (``torch.nn.Parameter``): A model parameter
      """
      if param.grad is not None:
          return param.grad
  
      # ZeRO stage 3 param
      if hasattr(param, "ds_id"):
          return param._z3_optimizer.get_fp32_grad_for_param(param)
  
      # ZeRO stage 1, 2, and bf16_optimizer params
      if hasattr(param, "_hp_mapping"):
          return param.get_full_hp_grad()
  
      return None
  
imports: import torch
  from dataclasses import dataclass
  from deepspeed import comm
  from typing import Dict
label_desc: Use image node_id 8 for calling a global function with example usage: safe_get_full_grad(param) and returns: None, param, param, param