code: def rotate_half(x):
      """Rotates half the hidden dims of the input."""
      x1 = x[..., : x.shape[-1] // 2]
      x2 = x[..., x.shape[-1] // 2 :]
      return torch.cat((-x2, x1), dim=-1)
  
file_path: ColossalAI/colossalai/inference/engine/modeling/llama.py
imports: import math
  from typing import List, Optional, Tuple
  import torch
  from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel
  from transformers.utils import logging
  from colossalai.inference.kv_cache.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import llama_context_attn_fwd, token_attention_fwd
  from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from ._utils import copy_kv_to_mem_cache
label_desc: Use image node_id 1 for calling a global function with example usage: rotate_half(x) and returns: torch