code: def forward(self, queries, keys, values, attn_mask):
      B, L, H, E = queries.shape
      _, S, _, D = values.shape
      scale = self.scale or 1.0 / sqrt(E)
  
      scores = torch.einsum("blhe,bshe->bhls", queries, keys)
  
      if self.mask_flag:
          if attn_mask is None:
              attn_mask = TriangularCausalMask(
                  B, L, device=queries.device
              )
  
          scores.masked_fill_(attn_mask.mask, -np.inf)
  
      A = self.dropout(torch.softmax(scale * scores, dim=-1))
      V = torch.einsum("bhls,bshd->blhd", A, values)
  
      if self.output_attention:
          return (V.contiguous(), A)
      else:
          return (V.contiguous(), None)
  
imports: import os
  import math
  import numpy
  from math import sqrt
  from scipy.fftpack import next_fast_len
label_desc: Use image node_id 2 for calling the FullAttention obj's underlying member method code with example usage: obj.forward(queries, keys, values, attn_mask) and returns: V,  A, V,  None