code: def make_sharded_optimizer_tensor(
      model_param: Union[ShardedTensor, ShardedTensorFactory],
      optim_param: torch.Tensor,
      prefix: str,
  ) -> Union[ShardedTensor, ShardedTensorFactory]:
      if isinstance(model_param, ShardedTensorFactory):
          return replace(
              model_param,
              key=f"{prefix}.{model_param.key}",
              data=optim_param,
          )
  
      assert (
          tuple(optim_param.shape) == model_param.local_shape
      ), f"Optimizer shape ({tuple(optim_param.shape)} does not match model shape ({model_param.local_shape})"
      return replace(
          model_param,
          key=f"{prefix}.{model_param.key}",
          data=optim_param,
          dtype=optim_param.dtype,
      )
  
file_path: megatron-lm/megatron/core/dist_checkpointing/optimizer.py
imports: import logging
  from copy import deepcopy
  from dataclasses import replace
  from itertools import chain
  from typing import Dict, Iterable, List, Tuple, Union
  import torch
  from .dict_utils import nested_values
  from .mapping import LocalNonpersitentObject, ShardedStateDict, ShardedTensor, ShardedTensorFactory, StateDict
  from .utils import extract_sharded_tensors, extract_sharded_tensors_and_factories
label_desc: Use image node_id 3 for calling a global function with example usage: make_sharded_optimizer_tensor(model_param, optim_param, prefix) and returns: replace, replace