code: def apply_rotary_embedding(q, k, cos, sin, index=None):
      """Helper function to apply Rotary Embeddings."""
      batch, qlen, qheads, d = q.shape
      kbatch, klen, kheads, kd = k.shape
      if index is not None:
          qcos = jax.lax.broadcast_in_dim(
              cos[index, :], (batch, qlen, qheads, d), (3,)
          )
          qsin = jax.lax.broadcast_in_dim(
              sin[index, :], (batch, qlen, qheads, d), (3,)
          )
      else:
          qcos = jax.lax.broadcast_in_dim(
              cos[:qlen, :], (batch, qlen, qheads, d), (1, 3)
          )
          qsin = jax.lax.broadcast_in_dim(
              sin[:qlen, :], (batch, qlen, qheads, d), (1, 3)
          )
      kcos = jax.lax.broadcast_in_dim(
          cos[:klen, :], (batch, klen, kheads, d), (1, 3)
      )
      ksin = jax.lax.broadcast_in_dim(
          sin[:klen, :], (batch, klen, kheads, d), (1, 3)
      )
      out_q = (q * qcos) + (rotate_half(q) * qsin)
      out_k = (k * kcos) + (rotate_half(k) * ksin)
      return out_q, out_k
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 6 for calling a global function with example usage: apply_rotary_embedding(q, k, cos, sin, index) and returns: out_q,  out_k