code: def forward(self, x, mask=None):
      b, n, _, h, device = *x.shape, self.heads, x.device
      softmax = torch.softmax
  
      qkv = self.to_qkv(x).chunk(3, dim=-1)
      q, k, v = map(
          lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv
      )
  
      q = q * self.scale
  
      dots = torch.einsum("b h i d, b h j d -> b h i j", q, k)
      mask_value = max_neg_value(dots)
  
      if exists(mask):
          mask = rearrange(mask, "b j -> b () () j")
          dots.masked_fill_(~mask, mask_value)
          del mask
  
      if self.causal:
          i, j = dots.shape[-2:]
          mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
          dots.masked_fill_(mask, mask_value)
  
      attn = softmax(dots, dim=-1)
  
      out = torch.einsum("b h i j, b h j d -> b h i d", attn, v)
      out = rearrange(out, "b h n d -> b n (h d)")
      out = self.to_out(out)
      return out
  
imports: import torch
  import torch.nn.functional
  from einops import rearrange
  from torch import nn
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, mask) and returns: out