code: def test_mha_even_sharding(n_heads: int, n_shards: int):
      """
      Even head sharding for MHA.
  
      Args:
          n_heads (int): The number QKV heads.
          n_shards (int): The number of shards to test for.
      """
      param = fill_with_head_ids(HEAD_SIZE, n_heads)
  
      n_local_heads = n_heads // n_shards
      sharded_shape = (HEAD_SIZE * n_heads, HEAD_SIZE * n_local_heads)
  
      for shard_rank in range(n_shards):
          sharded_param = shard_attn_out_param(
              param, shard_rank, n_shards, HEAD_SIZE
          )
          n_heads_local_q, _ = get_local_heads(
              shard_rank, n_shards, n_heads
          )
  
          assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
          assert sharded_param.shape == sharded_shape
  
          heads = torch.chunk(sharded_param, n_local_heads, dim=1)
  
          for i, head in enumerate(heads):
              assert torch.all(head == i + shard_rank * n_local_heads)
  
file_path: deepspeed/tests/unit/inference/v2/model_implementations/sharding/test_attn_out_sharding.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.sharding import *
label_desc: Use image node_id 2 for calling a global function with example usage: test_mha_even_sharding(n_heads, n_shards) without return types