code: def repeat_kv(
      hidden_states: torch.Tensor, n_rep: int
  ) -> torch.Tensor:
      """
      This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
      num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
      """
      batch, num_key_value_heads, slen, head_dim = hidden_states.shape
      if n_rep == 1:
          return hidden_states
      hidden_states = hidden_states[:, :, None, :, :].expand(
          batch, num_key_value_heads, n_rep, slen, head_dim
      )
      return hidden_states.reshape(
          batch, num_key_value_heads * n_rep, slen, head_dim
      )
  
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch import nn
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.llama.modeling_llama import LlamaConfig, LlamaRMSNorm
  from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from colossalai.kernel.cuda_native.mha.flash_attn_2 import HAS_FLASH_ATTN
  from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
  from colossalai.moe.layers import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.utils import get_activation, set_moe_args
label_desc: Use image node_id 8 for calling a global function with example usage: repeat_kv(hidden_states, n_rep) and returns: hidden_states, hidden_states