docstr: einstein notation
  b - batch
  h - heads
  n, i, j - sequence length (base sequence length, source, target)
  d - feature dimension
code: def forward(self, q, k, v, mask=None):
      """
      einstein notation
      b - batch
      h - heads
      n, i, j - sequence length (base sequence length, source, target)
      d - feature dimension
      """
  
      n, device = q.shape[-2], q.device
  
      scale = q.shape[-1] ** -0.5
  
      if self.use_flash:
          return self.flash_attn(q, k, v, mask=mask)
  
      kv_einsum_eq = "b j d" if k.ndim == 3 else "b h j d"
  
      # similarity
  
      sim = einsum(f"b h i d, {kv_einsum_eq} -> b h i j", q, k) * scale
  
      # key padding mask
  
      if exists(mask):
          mask = rearrange(mask, "b j -> b 1 1 j")
          sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)
  
      # causal mask
  
      if self.causal:
          causal_mask = self.get_mask(n, device)
          sim = sim.masked_fill(
              causal_mask, -torch.finfo(sim.dtype).max
          )
  
      # attention
  
      attn = sim.softmax(dim=-1)
      attn = self.attn_dropout(attn)
  
      # aggregate values
  
      out = einsum(f"b h i j, {kv_einsum_eq} -> b h i d", attn, v)
  
      return out
  
imports: from collections import namedtuple
  from functools import wraps
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
  from packaging import version
  from torch import einsum, nn
label_desc: Use image node_id 4 for calling the Attend obj's underlying member method code with example usage: obj.forward(q, k, v, mask) and returns: out, self