class_docstr: cross attention - using multi-query + one-headed key / values as in PaLM w/ optional parallel feedforward.
code: def __init__(
      self,
      dim,
      *,
      context_dim=None,
      dim_head=64,
      heads=8,
      parallel_ff=False,
      ff_mult=4,
      norm_context=False,
  ):
      super().__init__()
      self.heads = heads
      self.scale = dim_head**-0.5
      inner_dim = heads * dim_head
      if context_dim is None:
          context_dim = dim
  
      self.norm = LayerNorm(dim)
      self.context_norm = (
          LayerNorm(context_dim) if norm_context else nn.Identity()
      )
  
      self.to_q = nn.Linear(dim, inner_dim, bias=False)
      self.to_kv = nn.Linear(context_dim, dim_head * 2, bias=False)
      self.to_out = nn.Linear(inner_dim, dim, bias=False)
  
      # whether to have parallel feedforward
  
      ff_inner_dim = ff_mult * dim
  
      self.ff = (
          nn.Sequential(
              nn.Linear(dim, ff_inner_dim * 2, bias=False),
              SwiGLU(),
              nn.Linear(ff_inner_dim, dim, bias=False),
          )
          if parallel_ff
          else None
      )
  
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 1 to create a new CrossAttention object from inherited base classes: nn with example: obj = CrossAttention(dim)