code: def get_hp_fragment_mapping(
      lp_param,
      lp_start,
      flat_hp_partition,
      gradient_dict,
      offload_gradient_dict,
      use_offload,
      param_group_index,
      partition_start,
      partition_size,
      optimizer_state_dict,
  ):
      lp_end = lp_param.numel() + lp_start
      hp_start = partition_start
      hp_end = partition_start + partition_size
  
      fragment_start = max(lp_start, hp_start)
      fragment_end = min(lp_end, hp_end)
      assert (
          fragment_start < fragment_end
      ), f"fragment start {fragment_start} should be < fragment_end {fragment_end}"
  
      fragment_numel = fragment_end - fragment_start
      hp_frag_address = fragment_address(
          start=fragment_start - hp_start, numel=fragment_numel
      )
      hp_fragment_tensor = flat_hp_partition.narrow(
          0, hp_frag_address.start, hp_frag_address.numel
      )
      optim_fragment = {
          key: value.narrow(
              0, hp_frag_address.start, hp_frag_address.numel
          )
          for key, value in optimizer_state_dict.items()
          if torch.is_tensor(value)
          and value.shape == flat_hp_partition.shape
      }
  
      lp_frag_address = fragment_address(
          start=fragment_start - lp_start, numel=fragment_numel
      )
      lp_fragment_tensor = lp_param.flatten().narrow(
          0, lp_frag_address.start, lp_frag_address.numel
      )
  
      return tensor_fragment(
          lp_fragment=lp_fragment_tensor,
          lp_fragment_address=lp_frag_address,
          hp_fragment=hp_fragment_tensor,
          hp_fragment_address=hp_frag_address,
          optim_fragment=optim_fragment,
          gradient_dict=gradient_dict,
          offload_gradient_dict=offload_gradient_dict,
          use_offload=use_offload,
          param_group_index=param_group_index,
      )
  
imports: import torch
  from dataclasses import dataclass
  from deepspeed import comm
  from typing import Dict
label_desc: Use image node_id 14 for calling a global function with example usage: get_hp_fragment_mapping(lp_param, lp_start, flat_hp_partition, gradient_dict, offload_gradient_dict, use_offload, param_group_index, partition_start, partition_size, optimizer_state_dict) and returns: tensor_fragment