docstr: einops
  b - batch
  g - groups
  n - sequence
  d - feature dimension
docstr_tok: ["einops","b","-","batch","g","-","groups","n","-","sequence","d","-","feature","dimension"]
code: def forward(
      self,
      x,
      context=None,
      mask=None,
      queries_scale=None,
      keys_scale=None,
      values_scale=None,
      output_scale=None,
      rotary_emb: Optional[Tuple[Tensor, Tensor]] = None,
  ):
      """
      einops
      b - batch
      g - groups
      n - sequence
      d - feature dimension
      """
      b, g, h = x.shape[0], self.groups, self.heads
  
      one_expert = x.ndim == 3
  
      if one_expert:
          assert g == 1
          x = rearrange(x, "b n d -> b 1 n d")
  
      assert x.ndim == 4
      assert x.shape[1] == g
  
      # fold the groups into the feature dimension to be processed in one go by grouped convolutions
  
      x = rearrange(x, "b g n d -> b g d n")
  
      # handle context for cross attention
  
      if exists(context):
          context_one_expert = context.ndim == 3
  
          if context_one_expert:
              assert g == 1
              context = rearrange(context, "b n d -> b 1 n d")
  
          assert context.ndim == 4
          assert context.shape[1] == g
  
          context = rearrange(context, "b g n d -> b g d n")
  
      context = default(context, x)
  
      # take care of mask
  
      if exists(mask):
          if mask.ndim == 2:
              mask = repeat(mask, "b n -> (b g) n", g=g)
          elif mask.ndim == 3:
              mask = rearrange(mask, "b g n -> (b g) n")
  
          mask = F.pad(mask, (1, 0), value=True)
  
      # prenorm if applicable
  
      x = self.norm(x)
      context = self.context_norm(context)
  
      # fold groups into dimension for grouped conv
  
      x, context = map(
          lambda t: rearrange(t, "b g d n -> b (g d) n"), (x, context)
      )
  
      # queries, keys, values
  
      q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim=1))
  
      # split out heads and merge groups into batches
  
      q, k, v = map(
          lambda t: rearrange(t, "b (g h d) n -> b g h n d", h=h, g=g),
          (q, k, v),
      )
  
      # rotary embedding
  
      if exists(rotary_emb):
          q_rotary_emb, k_rotary_emb = rotary_emb
  
          if q_rotary_emb.ndim > 2:
              q_rotary_emb = rearrange(
                  q_rotary_emb, "b g n d -> b g 1 n d"
              )
  
          if k_rotary_emb.ndim > 2:
              k_rotary_emb = rearrange(
                  k_rotary_emb, "b g n d -> b g 1 n d"
              )
  
          q = apply_rotary_pos_emb(q_rotary_emb, q)
          k = apply_rotary_pos_emb(k_rotary_emb, k)
  
      # give gradients to routed keys / values via normalized scores from the router, if passed in
  
      if exists(queries_scale):
          q = q * queries_scale
  
      if exists(keys_scale):
          k = k * keys_scale
  
      if exists(values_scale):
          v = v * values_scale
  
      # merge group into batch
  
      q, k, v = map(
          lambda t: rearrange(t, "b g ... -> (b g) ..."), (q, k, v)
      )
  
      # concat null key / values, to protect against a row having all masked out elements and save a lot of headache
  
      nk, nv = map(
          lambda t: repeat(t, "g h 1 d -> (b g) h 1 d", b=b),
          self.null_kv,
      )
  
      k = torch.cat((nk, k), dim=-2)
      v = torch.cat((nv, v), dim=-2)
  
      # attention
  
      out = self.attend(q, k, v, mask=mask)
  
      # combine heads out
  
      out = rearrange(out, "(b g) h n d -> b (g h d) n", g=g)
  
      out = self.to_out(out)
  
      out = rearrange(out, "b (g d) n -> b g n d", g=g)
  
      if one_expert:
          out = rearrange(out, "b 1 n d -> b n d")
  
      if exists(output_scale):
          out = out * output_scale
  
      return out
  
imports: import math
  import torch
  import torch.nn.functional
  from torch import Tensor, nn, einsum
  from typing import Tuple, Optional
  from einops import rearrange, repeat, reduce, pack, unpack
  from mixture_of_attention.attend import Attend
  from mixture_of_attention.rotary_emb import apply_rotary_pos_emb
  from local_attention import LocalMHA
  from colt5_attention import CoordinateDescentRouter
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, context, mask, queries_scale, keys_scale, values_scale, output_scale, rotary_emb) and returns: out