code: def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
      gather_indices = position_ids[
          :, None, :, None
      ]  # [bs, 1, seq_len, 1]
      gather_indices = gather_indices.repeat(
          1, cos.shape[1], 1, cos.shape[3]
      )
      cos = torch.gather(
          cos.repeat(gather_indices.shape[0], 1, 1, 1).to(q.dtype),
          2,
          gather_indices,
      )
      sin = torch.gather(
          sin.repeat(gather_indices.shape[0], 1, 1, 1).to(k.dtype),
          2,
          gather_indices,
      )
      q_embed = (q * cos) + (rotate_half(q) * sin)
      k_embed = (k * cos) + (rotate_half(k) * sin)
      return q_embed, k_embed
  
file_path: long-lora/gptneox_attn_replace.py
imports: from typing import Optional, Tuple
  import warnings
  import torch
  import transformers
  from einops import rearrange
  from flash_attn import flash_attn_varlen_qkvpacked_func, flash_attn_varlen_func
  from flash_attn.bert_padding import unpad_input, pad_input
label_desc: Use image node_id 2 for calling a global function with example usage: apply_rotary_pos_emb(q, k, cos, sin, position_ids) and returns: q_embed,  k_embed