code: def load_qkv(
      config, prefix: str, weights, num_heads, head_size, hidden_size
  ):
      weight = weights.get_multi_weights_col(
          [prefix], quantize=config.quantize, dim=0
      )
      if isinstance(weight, torch.Tensor):
          # Only on non quantized versions
          weight = (
              weight.view(
                  num_heads,
                  3,
                  head_size,
                  hidden_size,
              )
              .permute(1, 0, 2, 3)
              .reshape(-1, hidden_size)
          )
  
      bias = weights.get_sharded(f"{prefix}.bias", dim=0)
      bias = (
          bias.view(num_heads, 3, head_size)
          .permute(1, 0, 2)
          .reshape(-1)
      )
  
      linear = get_linear(weight, bias, config.quantize)
      if config.use_parallel_residual:
          return linear
      else:
          return TensorParallelColumnLinear(linear)
  
imports: import torch
  import torch.distributed
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.gpt_neox import GPTNeoXConfig
  from typing import Optional, List, Tuple
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.flash_attn import attention
  from text_generation_server.utils.layers import TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelHead, FastLayerNorm, PositionRotaryEmbedding, get_linear
label_desc: Use image node_id 2 for calling a global function with example usage: load_qkv(config, prefix, weights, num_heads, head_size, hidden_size) and returns: linear, TensorParallelColumnLinear