code: def forward(self, x, context=None, mask=None, attn_mask=None):
      x = self.norm(x)
      kv_input = default(context, x)
  
      qkv = (self.to_q(x), *self.to_kv(kv_input).chunk(2, dim=-1))
  
      q, k, v = map(
          lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads),
          qkv,
      )
  
      q = self.q_norm(q)
      k = self.k_norm(k)
  
      dots = torch.matmul(q, k.transpose(-1, -2))
  
      if exists(mask):
          mask = rearrange(mask, "b j -> b 1 1 j")
          dots = dots.masked_fill(~mask, -torch.finfo(dots.dtype).max)
  
      if exists(attn_mask):
          dots = dots.masked_fill(
              ~attn_mask, -torch.finfo(dots.dtype).max
          )
  
      attn = self.attend(dots)
      attn = self.dropout(attn)
  
      out = torch.matmul(attn, v)
      out = rearrange(out, "b h n d -> b n (h d)")
      return self.to_out(out)
  
imports: from functools import partial
  from typing import List, Union
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from torch.nn.utils.rnn import pad_sequence
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, context, mask, attn_mask) and returns: self