code: def forward(self, x, cos, sin, position_ids):
      # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
      cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
      sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
      cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
      sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
      x_embed = (x * cos) + (rotate_half(x) * sin)
  
      return x_embed
  
imports: import math
  import os
  import types
  from collections import defaultdict
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch_int.nn.bmm import BMM_S8T_S8N_F32T, BMM_S8T_S8N_S8T
  from transformers import PreTrainedModel
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.configuration_llama import LlamaConfig
  from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, LlamaAttention, LlamaDecoderLayer, LlamaMLP, LlamaRotaryEmbedding, repeat_kv, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.tensor_parallel.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import copy_kv_cache_to_dest, int8_rotary_embedding_fwd, smooth_llama_context_attn_fwd, smooth_token_attention_fwd
  from .base_model import BaseSmoothForCausalLM
  from .linear import W8A8B8O8Linear, W8A8BFP32O32LinearSiLU, W8A8BFP32OFP32Linear
label_desc: Use image node_id 2 for calling the LlamaApplyRotary obj's underlying member method code with example usage: obj.forward(x, cos, sin, position_ids) and returns: x_embed