code: def make_tp_sharded_tensor_for_checkpoint(
      tensor,
      key,
      tp_axis=0,
      replica_id=None,
      prepend_offsets=(),
      **kwargs,
  ):
      """Helper for instantiating a ShardedTensor where the `tp_axis` dimension is sharded across TP group.
  
      Optionally, can provide offsets which prepend new dimensions to the tensor.
      """
  
      prepend_axis_num = len(prepend_offsets)
  
      if replica_id is None:
          replica_id = (
              0,
              0,
              parallel_state.get_data_parallel_rank(
                  with_context_parallel=True
              ),
          )
  
      return ShardedTensor.from_rank_offsets(
          key,
          tensor,
          *prepend_offsets,
          (
              tp_axis + prepend_axis_num,
              parallel_state.get_tensor_model_parallel_rank(),
              parallel_state.get_tensor_model_parallel_world_size(),
          ),
          replica_id=replica_id,
          prepend_axis_num=prepend_axis_num,
          **kwargs,
      )
  
imports: import math
  import operator
  from functools import reduce
  import torch
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.mapping import ShardedTensor
label_desc: Use image node_id 12 for calling a global function with example usage: make_tp_sharded_tensor_for_checkpoint(tensor, key, tp_axis, replica_id, prepend_offsets) and returns: ShardedTensor