class_docstr: cross attention - using multi-query + one-headed key / values as in PaLM w/ optional parallel feedforward.
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, context):
      """
      einstein notation
      b - batch
      h - heads
      n, i, j - sequence length (base sequence length, source, target)
      d - feature dimension
      """
  
      # pre-layernorm, for queries and context
  
      x = self.norm(x)
      context = self.context_norm(context)
  
      # get queries
  
      q = self.to_q(x)
      q = rearrange(q, "b n (h d) -> b h n d", h=self.heads)
  
      # scale
  
      q = q * self.scale
  
      # get key / values
  
      k, v = self.to_kv(context).chunk(2, dim=-1)
  
      # query / key similarity
  
      sim = einsum("b h i d, b j d -> b h i j", q, k)
  
      # attention
  
      sim = sim - sim.amax(dim=-1, keepdim=True)
      attn = sim.softmax(dim=-1)
  
      # aggregate
  
      out = einsum("b h i j, b j d -> b h i d", attn, v)
  
      # merge and combine heads
  
      out = rearrange(out, "b h n d -> b n (h d)")
      out = self.to_out(out)
  
      # add parallel feedforward (for multimodal layers)
  
      if self.ff is not None:
          # pylint: disable=E1102
          out = out + self.ff(x)
  
      return out
  
imports: import torch
  from torch import einsum, nn
  import torch.nn.functional
  from einops import rearrange
  from towhee.models.layers.activations.swiglu import SwiGLU
label_desc: Use image node_id 2 for calling the CrossAttention obj's underlying member method code with example usage: obj.forward(x, context) and returns: out