docstr: einstein notation
  b - batch
  h - heads
  n, i, j - sequence length (base sequence length, source, target)
  d - feature dimension
docstr_tok: ["einstein","notation","b","-","batch","h","-","heads","n",",","i",",","j","-","sequence","length","(","base","sequence","length",",","source",",","target",")","d","-","feature","dimension"]
code: def forward(self, x):
      """
      einstein notation
      b - batch
      h - heads
      n, i, j - sequence length (base sequence length, source, target)
      d - feature dimension
      """
  
      n, device, h = x.shape[1], x.device, self.heads
  
      # pre layernorm
  
      x = self.norm(x)
  
      # queries, keys, values
  
      q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim=-1))
  
      # split heads
      # they use multi-query single-key-value attention, yet another Noam Shazeer paper
      # they found no performance loss past a certain scale, and more efficient decoding obviously
      # https://arxiv.org/abs/1911.02150
  
      q = rearrange(q, "b n (h d) -> b h n d", h=h)
  
      # rotary embeddings
  
      positions = self.get_rotary_embedding(n, device)
      q, k = map(lambda t: apply_rotary_pos_emb(positions, t), (q, k))
  
      # scale
  
      q = q * self.scale
  
      b, h, i, d, j = (
          q.size(0),
          q.size(1),
          q.size(2),
          q.size(3),
          k.size(1),
      )
  
      # similarity
  
      # sim = einsum("b h i d, b j d -> b h i j", q, k)
      sim = matmul(q.reshape(b, h * i, d), k.transpose(1, 2))
      sim = sim.reshape(b, h, i, j)
  
      # causal mask
  
      causal_mask = self.get_mask(n, device)
      sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
  
      # attention
  
      sim = sim - sim.amax(dim=-1, keepdim=True).detach()
      attn = sim.softmax(dim=-1)
  
      b_, h_, i_, j_, d_ = (
          attn.size(0),
          attn.size(1),
          attn.size(2),
          attn.size(3),
          v.size(2),
      )
  
      # aggregate values
  
      # out = einsum("b h i j, b j d -> b h i d", attn, v)
      out = matmul(attn.reshape(b_, h_ * i_, j_), v)
      out = out.reshape(b_, h_, i_, d_)
  
      # merge heads
  
      out = rearrange(out, "b h n d -> b n (h d)")
      return self.to_out(out)
  
imports: import torch
  import torch.nn.functional
  from einops import rearrange
  from torch import matmul, nn
label_desc: Use image node_id 4 for calling the Attention obj's underlying member method code with example usage: obj.forward(x) and returns: self