code: def shard_unembed_param(
      param: torch.Tensor, shard_rank: int, num_shards: int
  ) -> torch.Tensor:
      """
      Utility method for sharding an unembed parameter. We shard unembeddings on the vocab dimension
      with the expectation of an all-gather to produce the full results.
  
      TODO(cmikeh2): Really ideal would be if MII could have access to the comm and we would do
      an A2A and sharded sampling.
  
      Args:
          param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
          shard_rank (int): Which shard of the partitioned tensor to return.
          num_shards (int): The total number of shards the parameter is distributed across.
  
      Returns:
          torch.Tensor: The sharded parameter of shape [sharded_vocab_size, model_dim]
      """
      return shard_param(
          param,
          ShardingType.OUTER_DIMENSION,
          shard_rank,
          num_shards,
          granularity=1,
      )
  
imports: import torch
  from .types import ShardingType
  from .utils import shard_param, get_shard_endpoints
label_desc: Use image node_id 1 for calling a global function with example usage: shard_unembed_param(param, shard_rank, num_shards) and returns: shard_param