code: def test_mha_input_shape_error():
      param = torch.empty(256, 128)
  
      n_heads = 2
      head_size = 64
  
      with pytest.raises(ValueError):
          shard_qkv_param(param, 0, 1, 64)
  
file_path: deepspeed/tests/unit/inference/v2/model_implementations/sharding/test_qkv_sharding.py
imports: from typing import Optional
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.sharding import *
label_desc: Use image node_id 9 for calling a global function with example usage: test_mha_input_shape_error() without return types