code: def apply_rotary_pos_emb(
      q, k, cos, sin, position_ids, unsqueeze_dim=1
  ):
      """Applies Rotary Position Embedding to the query and key tensors.
  
      Args:
          q (`torch.Tensor`): The query tensor.
          k (`torch.Tensor`): The key tensor.
          cos (`torch.Tensor`): The cosine part of the rotary embedding.
          sin (`torch.Tensor`): The sine part of the rotary embedding.
          position_ids (`torch.Tensor`):
              The position indices of the tokens corresponding to the query and key tensors. For example, this can be
              used to pass offsetted position ids when working with a KV-cache.
          unsqueeze_dim (`int`, *optional*, defaults to 1):
              The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
              sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
              that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
              k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
              cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
              the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
      Returns:
          `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
      """
      cos = cos[position_ids].unsqueeze(unsqueeze_dim)
      sin = sin[position_ids].unsqueeze(unsqueeze_dim)
      q_embed = (q * cos) + (rotate_half(q) * sin)
      k_embed = (k * cos) + (rotate_half(k) * sin)
      return q_embed, k_embed
  
imports: import math
  from typing import Optional
  import pytest
  import torch
  import torch.utils.checkpoint
  from helpers import assert_similar, random_tensor
  from torch import nn
  from quanto import Calibration, quantize
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, unsqueeze_dim) and returns: q_embed,  k_embed