code: def get_file_indices_for_rank(self, pp_index, tp_index, dp_index):
      assert dp_index < len(
          self._3d_file_map
      ), f"DP index {dp_index} >= DP degree {len(self._3d_file_map)}"
      dp_2d_map = self._3d_file_map[dp_index]
      return dp_2d_map.get_data(pp_index, tp_index)
  
imports: import torch
  from .constants import BASE_OPTIMIZER_STATE, GROUP_PADDINGS, OPTIMIZER_STATE_DICT, PARTITION_COUNT
  from .reshape_utils import basic_folder_validation, get_zero_files, merge_state
  from .reshape_3d_utils import model_3d_desc, get_model_3d_descriptor
label_desc: Use image node_id 6 for calling the ZeROCheckpoint obj's underlying member method code with example usage: obj.get_file_indices_for_rank(pp_index, tp_index, dp_index) and returns: dp_2d_map