code: def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
      # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
      # pylint: disable=invalid-name
      cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
      sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
      cos = cos[position_ids].unsqueeze(
          1
      )  # [batch_size, 1, seq_len, dim]
      sin = sin[position_ids].unsqueeze(
          1
      )  # [batch_size, 1, seq_len, dim]
      q_embed = (q * cos) + (rotate_half(q) * sin)
      k_embed = (k * cos) + (rotate_half(k) * sin)
      return q_embed, k_embed
  
file_path: axolotl/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
imports: import importlib
  import math
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from accelerate import init_empty_weights
  from einops import rearrange
  from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
  from torch import nn
  from transformers import AutoConfig, AutoModelForCausalLM
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.utils import logging
  from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
label_desc: Use image node_id 3 for calling a global function with example usage: apply_rotary_pos_emb(q, k, cos, sin, position_ids) and returns: q_embed,  k_embed