code: def megatron_qkv_reshape(
      param: torch.Tensor, head_size: int, n_heads: int
  ) -> torch.Tensor:
      assert param.shape[0] == 3 * n_heads * head_size
  
      all_heads = torch.chunk(param, chunks=3 * n_heads, dim=0)
      q_heads = all_heads[::3]
      k_heads = all_heads[1::3]
      v_heads = all_heads[2::3]
      return torch.cat([q_heads, k_heads, v_heads], dim=0)
  
file_path: deepspeed/deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py
imports: import torch
  from ...model_implementations.parameter_base import ParameterBase
label_desc: Use image node_id 1 for calling a global function with example usage: megatron_qkv_reshape(param, head_size, n_heads) and returns: torch