code: def rotate(x):
      """Rotate function."""
      _, l, d = x.shape
      inv_freq = jnp.exp(jnp.arange(0, d, 2) * -(jnp.log(10000.0) / d))
      positions = jnp.arange(l)
      freqs = jnp.einsum("i,j->ij", positions, inv_freq)
      emb = jnp.concatenate((freqs, freqs), axis=-1)
      cos = jnp.cos(emb)
      sin = jnp.sin(emb)
  
      def mul(vecs, pos_emb):
          return jnp.einsum("bld,ld->bld", vecs, pos_emb)
  
      def rotate_half(x):
          x1, x2 = (
              x[..., : x.shape[-1] // 2],
              x[..., x.shape[-1] // 2 :],
          )
          return jnp.concatenate((-x2, x1), axis=x1.ndim - 1)
  
      return mul(x, cos) + mul(rotate_half(x), sin)
  
imports: from trax.fastmath import numpy
  from trax.layers import core
label_desc: Use image node_id 1 for calling a global function with example usage: rotate(x) and returns: unknown, jnp, jnp