code: def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
      """
      x: (batch_size, seqlen, nheads, headdim)
      cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
      """
      ro_dim = cos.shape[-1] * 2
      assert ro_dim <= x.shape[-1]
      cos = repeat(
          cos,
          "... d -> ... 1 (2 d)"
          if not interleaved
          else "... d -> ... 1 (d 2)",
      )
      sin = repeat(
          sin,
          "... d -> ... 1 (2 d)"
          if not interleaved
          else "... d -> ... 1 (d 2)",
      )
      return torch.cat(
          [
              x[..., :ro_dim] * cos
              + rotate_half(x[..., :ro_dim], interleaved) * sin,
              x[..., ro_dim:],
          ],
          dim=-1,
      )
  
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  from einops import rearrange, repeat
  from flash_attn.ops.triton.rotary import apply_rotary
label_desc: Use image node_id 2 for calling a global function with example usage: apply_rotary_emb_torch(x, cos, sin, interleaved) and returns: torch