code: def shard_attn_out_param(
      param: torch.Tensor,
      shard_rank: int,
      num_shards: int,
      head_size: int,
      n_heads_q: Optional[int] = None,
      n_heads_kv: Optional[int] = None,
  ) -> Optional[torch.Tensor]:
      """
      Utility method for sharding an attention output parameter.
      """
      if len(param.shape) == 1:
          # We will do the bias addition on the 0th rank only rather than scale the parameter and
          # implicitly reconstruct this in the distributed reduce.
          return param if shard_rank == 0 else None
  
      assert n_heads_kv is None or (
          n_heads_q is not None and n_heads_kv is not None
      ), "n_heads_kv should not be passed without n_heads_q"
  
      mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
  
      if mha_sharding:
          return shard_param(
              param,
              ShardingType.INNER_DIMENSION,
              shard_rank,
              num_shards,
              granularity=head_size,
          )
      else:
          assert (
              param.shape[0] == head_size * n_heads_q
          ), "GQA param shape is not correct"
  
          # 32 KV heads, 16 shards for example
          even_kv_sharding = n_heads_kv % num_shards == 0
  
          # 8 KV heads, 16 shards for example
          even_kv_distribution = num_shards % n_heads_kv == 0
  
          assert (
              even_kv_sharding or even_kv_distribution
          ), "No partitioning algorithm for this yet."
  
          if even_kv_sharding:
              # Same as original sharding scenario
              return shard_param(
                  param,
                  ShardingType.INNER_DIMENSION,
                  shard_rank,
                  num_shards,
                  granularity=head_size,
              )
          else:
              # We will first do a sharding on the KV and Q to map to the one KV shard per group of Q.
              q_sharding_degree = num_shards // n_heads_kv
  
              kv_head = shard_rank // q_sharding_degree
  
              q_sharding_rank = shard_rank % q_sharding_degree
              q_factor = n_heads_q // n_heads_kv
  
              q_chunk = param[
                  ...,
                  q_factor
                  * kv_head
                  * head_size : q_factor
                  * (kv_head + 1)
                  * head_size,
              ]
  
              return shard_param(
                  q_chunk,
                  ShardingType.INNER_DIMENSION,
                  q_sharding_rank,
                  q_sharding_degree,
                  granularity=head_size,
              )
  
imports: from typing import Optional
  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_attn_out_param(param, shard_rank, num_shards, head_size, n_heads_q, n_heads_kv) and returns: unknown, shard_param, shard_param, shard_param