code: def shard_param(
      param: Optional[torch.Tensor],
      shard_mode: ShardingType,
      shard_rank: int,
      num_shards: int,
      num_concatenated_matrices: int = 1,
      granularity: int = 32,
      bias_dims: int = 1,
  ) -> torch.Tensor:
      """
      Utility for sharding a parameter. This will return the slice of the parameter that should
      exist on the given shard_rank given the sharding configuration. The workflow here is
      to find the minimum bounded Tensor to shard, get the slicing endpoints, and then concatenate
      as needed.
  
      The typical use of this is as an internal helper function, so see if there is a higher level
      API that better suits the application.
  
      Args:
          param (torch.Tensor): The parameter to shard.
          shard_mode (ShardingType): The type of sharding to apply. See ShardingType for more context.
          shard_rank (int): The rank of the shard to return.
          num_shards (int): Total number of shards the parameter will be distrbuted across.
          num_concatenated_matrices (int): The number of matrices that have been concatenated together in the original
              parameter. An example of this is a fused QKV projection matrix, where the `num_concatenated_matrices`
              argument would be 3.
          granularity (int): The minimum alignment of the shard endpoints. For attention projection matrices, this
              should be set to the head size to support non-even sharding.
          bias_dims (int): The number of dimensions that are considered bias dimensions. This is used to support
              sharding of MoE and non-MoE biases on the same codepath.
      """
      assert (
          shard_rank < num_shards
      ), "Shard rank must be less than num_shards"
  
      # Easier to hide this inside of the sharding logic than to add checks in every model
      # implementation.
      if param is None:
          return None
  
      if num_shards == 1:
          # Trivial case of no sharding.
          return param
  
      if shard_mode == ShardingType.OUTER_DIMENSION:
  
          def get_matrices(dim_idx: int) -> torch.Tensor:
              dim_size = (
                  param.size(dim_idx) // num_concatenated_matrices
              )
              start_channel_id, end_channel_id = get_shard_endpoints(
                  dim_size, shard_rank, num_shards, granularity
              )
              return (
                  torch.chunk(
                      param, num_concatenated_matrices, dim=dim_idx
                  ),
                  start_channel_id,
                  end_channel_id,
              )
  
          if param.ndim == bias_dims:
              # Special case for bias parameters.
              matrices, start_channel_id, end_channel_id = get_matrices(
                  dim_idx=-1
              )
              return torch.cat(
                  [
                      mat[..., start_channel_id:end_channel_id]
                      for mat in matrices
                  ],
                  dim=-1,
              )
          else:
              # General case for weight parameters. This assumes MoE parameters are stored in the format of
              # [num_experts, out_features, in_features]
              matrices, start_channel_id, end_channel_id = get_matrices(
                  dim_idx=-2
              )
              return torch.cat(
                  [
                      mat[..., start_channel_id:end_channel_id, :]
                      for mat in matrices
                  ],
                  dim=-2,
              )
  
      elif shard_mode == ShardingType.INNER_DIMENSION:
          dim_size = param.size(-1) // num_concatenated_matrices
          start_channel_id, end_channel_id = get_shard_endpoints(
              dim_size, shard_rank, num_shards, granularity
          )
          matrices = torch.chunk(
              param, num_concatenated_matrices, dim=-1
          )
          return torch.cat(
              [
                  mat[..., start_channel_id:end_channel_id]
                  for mat in matrices
              ],
              dim=-1,
          )
  
imports: from typing import Optional, Tuple
  import torch
  from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
label_desc: Use image node_id 2 for calling a global function with example usage: shard_param(param, shard_mode, shard_rank, num_shards, num_concatenated_matrices, granularity, bias_dims) and returns: None, param, torch,  start_channel_id,  end_channel_id, torch, torch, torch