class_docstr: Parallel attention and feedforward with residual discovered by Wang et al + EleutherAI from GPT-J fame.
  
  Args:
      dim (`int`):
          Dimension of features.
      dim_head (`int`):
          Head of fused dimensions.
      heads (`int`):
          Number of heads.
      ff_mult (`int`):
          Multiple number for feedforward.
class_docstr_tok: ["Parallel","attention","and","feedforward","with","residual","discovered","by","Wang","et","al","+","EleutherAI","from","GPT-J","fame",".","Args",":","dim","(","`","int","`",")",":","Dimension","of","features",".","dim_head","(","`","int","`",")",":","Head","of","fused","dimensions",".","heads","(","`","int","`",")",":","Number","of","heads",".","ff_mult","(","`","int","`",")",":","Multiple","number","for","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, attn_mask=None):
      """
      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)
  
      # attention queries, keys, values, and feedforward inner
      q, k, v, ff = self.fused_attn_ff_proj(x).split(
          self.fused_dims, 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
  
      # similarity
      sim = einsum("b h i d, b j d -> b h i j", q, k)
  
      # causal mask
      causal_mask = self.get_mask(n, device)
      sim = sim.masked_fill(causal_mask, -torch.finfo(sim.dtype).max)
  
      # extra attention mask - for masking out attention from text CLS token to padding
      if attn_mask is not None:
          attn_mask = rearrange(attn_mask, "b i j -> b 1 i j")
          # pylint: disable=E1130
          sim = sim.masked_fill(~attn_mask, -torch.finfo(sim.dtype).max)
  
      # attention
      sim = sim - sim.amax(dim=-1, keepdim=True).detach()
      attn = sim.softmax(dim=-1)
  
      # aggregate values
      out = einsum("b h i j, b j d -> b h i d", attn, v)
  
      # merge heads
      out = rearrange(out, "b h n d -> b n (h d)")
      return self.attn_out(out) + self.ff_out(ff)
  
file_path: towhee/towhee/models/coca/coca.py
imports: import torch
  from torch import einsum, nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from towhee.models.layers.cross_attention import LayerNorm, Residual, RotaryEmbedding, apply_rotary_pos_emb, CrossAttention
  from towhee.models.layers.activations.swiglu import SwiGLU
label_desc: Use image node_id 4 for calling the ParallelTransformerBlock obj's underlying member method code with example usage: obj.forward(x, attn_mask) and returns: unknown