code: def self_attn(self, q, v, k, attn_bias):
      if self.is_rotary:
          k_rot = k[:, :, : self.pe_rotary_dims]
          k_pass = k[:, :, self.pe_rotary_dims :]
  
          q_rot = q[:, :, : self.pe_rotary_dims]
          q_pass = q[:, :, self.pe_rotary_dims :]
  
          sincos = fixed_pos_embedding(k_rot)
          q_rot = apply_rotary_pos_emb(q_rot, sincos)
          k_rot = apply_rotary_pos_emb(k_rot, sincos)
  
          k = jnp.concatenate([k_rot, k_pass], axis=-1)
          q = jnp.concatenate([q_rot, q_pass], axis=-1)
  
      attention_logits = jnp.einsum("thd,Thd->htT", q, k)
  
      sqrt_key_size = np.sqrt(self.dim_per_head).astype(k.dtype)
      attention_logits = attention_logits / sqrt_key_size
  
      attention_logits += attn_bias
  
      attention_weights = jax.nn.softmax(attention_logits)
      attention_vec = jnp.einsum(
          "htT,Thd->thd", attention_weights, v
      ).reshape((-1, self.dim_per_shard))
  
      return self.o(attention_vec)
  
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 2 for calling the TransformerLayerShard obj's underlying member method code with example usage: obj.self_attn(q, v, k, attn_bias) and returns: self