code: def apply_rotary_pos_emb(x, sincos):
      sin, cos = map(
          lambda t: repeat(t, "b n -> b (n j)", j=2)[
              -x.shape[0] :, None, :
          ],
          sincos,
      )
      return (x * cos) + (rotate_every_two(x) * sin)
  
imports: import haiku
  import jax
  import jax.numpy
  import numpy
  from einops import rearrange, repeat
  from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
  from jax.experimental import PartitionSpec
  from jax.experimental.maps import thread_resources
label_desc: Use image node_id 4 for calling a global function with example usage: apply_rotary_pos_emb(x, sincos) and returns: unknown