code: def apply_rotary_emb(xq, xk, freqs_cis) -> Tuple[Tensor, Tensor]:
      assert (
          freqs_cis.shape[1] == xq.shape[1] == xk.shape[1]
      ), f"freqs_cis shape mismatch {freqs_cis.shape} xq:{xq.shape} xk:{xk.shape}"
      xq = xq.reshape(*xq.shape[0:-1], -1, 2)
      xk = xk.reshape(*xk.shape[0:-1], -1, 2)
      assert len(xq.shape) == len(xk.shape) == len(freqs_cis.shape) == 5
      c, d = freqs_cis[..., 0:1], freqs_cis[..., 1:2]
      xq_out = complex_mult(xq, c, d)
      xk_out = complex_mult(xk, c, d)
      return xq_out.flatten(3), xk_out.flatten(3)
  
imports: from typing import Tuple, Union, Optional, Dict
  from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
  from tinygrad.helpers import getenv
label_desc: Use image node_id 3 for calling a global function with example usage: apply_rotary_emb(xq, xk, freqs_cis) and returns: xq_out,  xk_out