docstr: ..NOTE: We're reusing the einsum notation suggested by the paper, changed in that
  heads are folded in the batch dimension
code: def forward(
      self,
      q: torch.Tensor,
      k: torch.Tensor,
      v: torch.Tensor,
      *args,
      **kwargs,
  ):
      """..NOTE: We're reusing the einsum notation suggested by the paper, changed in that
      heads are folded in the batch dimension"""
  
      content_lambda = torch.einsum(
          "bnk,bnv->bkv", torch.softmax(k, dim=-1), v
      )
      content_output = torch.einsum("bnk,bkv->bnv", q, content_lambda)
  
      rel_pos_emb = self.rel_pos_emb[self.rel_pos]
  
      # Handle real sequence length being possibly smaller
      seq_len = q.shape[1]
      rel_pos_emb = rel_pos_emb[:seq_len, :seq_len, :]
  
      # Compute the position lambda for every possible combination in one go, then compute the
      # position related contribution
      position_lambdas = torch.einsum(
          "mnk,bnv->bnkv", rel_pos_emb, v
      )  # one lambda per position
      position_output = (q.unsqueeze(2) @ position_lambdas).squeeze()
      att = content_output + position_output
  
      att = self.attn_drop(att)
  
      return att
  
file_path: xformers/xformers/components/attention/lambda_layer.py
imports: from dataclasses import dataclass
  import torch
  from xformers.components.attention import Attention, AttentionConfig, register_attention
label_desc: Use image node_id 2 for calling the LambdaLayer obj's underlying member method code with example usage: obj.forward(q, k, v) and returns: att