code: def apply_rotary_emb(query, key, rope, interleave):
      if interleave:
          query = query.transpose(1, 2)
          key = key.transpose(1, 2)
          query_ = query.float().reshape(*query.shape[:-1], -1, 2)
          query_ = torch.view_as_complex(query_)
          key_ = key.float().reshape(*key.shape[:-1], -1, 2)
          key_ = torch.view_as_complex(key_)
          rope = rope[:, : rope.size(1) // 2].view(
              1, query_.size(1), 1, query_.size(3)
          )
          query_out = torch.view_as_real(query_ * rope).flatten(3)
          key_out = torch.view_as_real(key_ * rope).flatten(3)
          return query_out.transpose(1, 2).type_as(
              query
          ), key_out.transpose(1, 2).type_as(key)
      else:
          cos, sin = rope.real, rope.imag
          rotary_dim = cos.size(1)
          head_dim = query.size(3)
          if rotary_dim < head_dim:
              q_embed = (query[:, :, :, :rotary_dim] * cos) + (
                  rotate_half(query[:, :, :, :rotary_dim]) * sin
              )
              k_embed = (key[:, :, :, :rotary_dim] * cos) + (
                  rotate_half(key[:, :, :, :rotary_dim]) * sin
              )
              q_embed = torch.cat(
                  [q_embed, query[:, :, :, rotary_dim:]], dim=-1
              )
              k_embed = torch.cat(
                  [k_embed, key[:, :, :, rotary_dim:]], dim=-1
              )
          else:
              q_embed = (query * cos) + (rotate_half(query) * sin)
              k_embed = (key * cos) + (rotate_half(key) * sin)
          return q_embed.type_as(query), k_embed.type_as(key)
  
imports: import torch
  import torch.nn
  from math import log, sqrt
  from torch import Tensor
  from typing import Optional, Tuple
  from torch.nn.functional import scaled_dot_product_attention
  from torch.utils.checkpoint import checkpoint
  from torch.nn.utils import skip_init
  from .alibi_position_bias import AlibiPositionalBias
  from torch.distributed import all_reduce
  from importlib import import_module
label_desc: Use image node_id 3 for calling a global function with example usage: apply_rotary_emb(query, key, rope, interleave) and returns: query_out,  key_out, q_embed,  k_embed