class_docstr: A module which performs QKV attention.
  Attention mask is accepted as input.
docstr: Apply QKV attention with attention mask.
  
  Args:
      q: an [N x d x n_seq1] of queries.
      k: an [N x d x n_seq2] of keys.
      v: an [N x d x n_seq2] of values.
      mask: Attention mask of size N x n_seq1 x n_seq2
  
  Returns: an [N x d x n_seq1] tensor after attention.
docstr_tok: ["Apply","QKV","attention","with","attention","mask",".","Args",":","q",":","an","[","N","x","d","x","n_seq1","]","of","queries",".","k",":","an","[","N","x","d","x","n_seq2","]","of","keys",".","v",":","an","[","N","x","d","x","n_seq2","]","of","values",".","mask",":","Attention","mask","of","size","N","x","n_seq1","x","n_seq2","Returns",":","an","[","N","x","d","x","n_seq1","]","tensor","after","attention","."]
code: def forward(self, q, k, v, mask):
      r"""
      Apply QKV attention with attention mask.
  
      Args:
          q: an [N x d x n_seq1] of queries.
          k: an [N x d x n_seq2] of keys.
          v: an [N x d x n_seq2] of values.
          mask: Attention mask of size N x n_seq1 x n_seq2
  
      Returns: an [N x d x n_seq1] tensor after attention.
      """
  
      bs, width, length_q = q.shape
      _, _, length_k = k.shape
  
      assert width % self.n_heads == 0
      ch = width // self.n_heads
  
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = torch.einsum(
          "bct,bcs->bts",
          (q * scale).view(bs * self.n_heads, ch, length_q),
          (k * scale).view(bs * self.n_heads, ch, length_k),
      )  # More stable with f16 than dividing afterwards
  
      # Duplicate mask n_heads times
      mask = mask.repeat_interleave(self.n_heads, dim=0)
      assert mask.shape == weight.shape
      max_neg_value = -float("inf")
      weight = weight.masked_fill(~mask, max_neg_value)
  
      weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
  
      # When we use an arbitrary mask, there is a possibility that we get nans in softmax.
      # In this case, use nan_to_num to make it a non-nan number.
      weight = weight.nan_to_num_()
      a = torch.einsum(
          "bts,bcs->bct",
          weight,
          v.reshape(bs * self.n_heads, ch, length_k),
      )
      # We also return weight here for attention visualization.
      return a.reshape(bs, -1, length_q), weight
  
imports: import math
  import numpy
  import torch
  import torch.nn
  from torch.cuda.amp import custom_bwd, custom_fwd
label_desc: Use image node_id 2 for calling the QKVMaskedAttention obj's underlying member method code with example usage: obj.forward(q, k, v, mask) and returns: a,  weight