code: def shard_embedding_param(
      param: torch.Tensor, shard_rank: int, num_shards: int
  ) -> torch.Tensor:
      """
      Utility method for sharding an embedding parameter.
  
      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.
      """
      return shard_param(
          param, ShardingType.INNER_DIMENSION, shard_rank, num_shards
      )
  
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_embedding_param(param, shard_rank, num_shards) and returns: shard_param