code: def _get_extra_state_offsets(
      sharded_offsets: Iterable[Tuple[int, int, int]]
  ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
      """Turns ShardedTensor offsets into offsets suitable for ShardedObject."""
      if sharded_offsets:
          sharded_offsets = sorted(
              sharded_offsets, key=itemgetter(0)
          )  # sort by axis
          axis, extra_state_offset, extra_state_shape = zip(
              *sharded_offsets
          )
          assert list(axis) == list(
              range(len(axis))
          ), f"Expected contiguous axis for offsets: {sharded_offsets}"
      else:
          extra_state_shape = (1,)
          extra_state_offset = (0,)
      return extra_state_shape, extra_state_offset
  
file_path: megatron-lm/megatron/core/transformer/utils.py
imports: from operator import itemgetter
  from typing import Any, Dict, Iterable, Optional, Tuple, Union
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedObject, StateDict
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 8 for calling a global function with example usage: _get_extra_state_offsets(sharded_offsets) and returns: extra_state_shape,  extra_state_offset