code: def make_sharded_tensor_for_checkpoint(
      tensor, key, prepend_offsets=(), replica_id=None, **kwargs
  ):
      """Helper for instantiating a non-sharded ShardedTensor (replicated across TP and DP 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,
              parallel_state.get_tensor_model_parallel_rank(),
              parallel_state.get_data_parallel_rank(
                  with_context_parallel=True
              ),
          )
  
      return ShardedTensor.from_rank_offsets(
          key,
          tensor,
          *prepend_offsets,
          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 13 for calling a global function with example usage: make_sharded_tensor_for_checkpoint(tensor, key, prepend_offsets, replica_id) and returns: ShardedTensor