code: def test_gqa_input_shape_error():
      head_size = 64
      n_heads_q = 16
      n_heads_kv = 4
  
      # Correct shape is 1536 (=16 * 64 + 2 * 4 * 64), 1024
      param = torch.empty(2048, 1024)
  
      with pytest.raises(ValueError):
          shard_qkv_param(param, 0, 1, head_size, n_heads_q, n_heads_kv)
  
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 10 for calling a global function with example usage: test_gqa_input_shape_error() without return types