code: def rotary_forward(q, k, cos, sin, position_ids):
      cos = cos[position_ids].unsqueeze(1)
      sin = sin[position_ids].unsqueeze(1)
  
      chunk_size = q.shape[-1] // 2
      q1, q2 = q.split(chunk_size, -1)
      q_rotated = torch.cat((-q2, q1), dim=-1)
      k1, k2 = k.split(chunk_size, -1)
      k_rotated = torch.cat((-k2, k1), dim=-1)
  
      q_embed = (q * cos) + (q_rotated * sin)
      k_embed = (k * cos) + (k_rotated * sin)
      return q_embed, k_embed
  
imports: from typing import Optional, Tuple, Union
  import os
  import torch
  import torch.distributed
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.activations import ACT2FN
  from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers import GPTNeoXConfig
  from loguru import logger
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 4 for calling a global function with example usage: rotary_forward(q, k, cos, sin, position_ids) and returns: q_embed,  k_embed