code: def forward(
      self,
      x,
      context=None,
      mask=None,
      additional_tokens=None,
      n_times_crossframe_attn_in_self=0,
  ):
      h = self.heads
  
      if additional_tokens is not None:
          # get the number of masked tokens at the beginning of the output sequence
          n_tokens_to_mask = additional_tokens.shape[1]
          # add additional token
          x = torch.cat([additional_tokens, x], dim=1)
  
      q = self.to_q(x)
      context = default(context, x)
      k = self.to_k(context)
      v = self.to_v(context)
  
      if n_times_crossframe_attn_in_self:
          # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439
          assert x.shape[0] % n_times_crossframe_attn_in_self == 0
          n_cp = x.shape[0] // n_times_crossframe_attn_in_self
          k = repeat(
              k[::n_times_crossframe_attn_in_self],
              "b ... -> (b n) ...",
              n=n_cp,
          )
          v = repeat(
              v[::n_times_crossframe_attn_in_self],
              "b ... -> (b n) ...",
              n=n_cp,
          )
  
      q, k, v = map(
          lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)
      )
  
      ## old
      """
          sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
          del q, k
  
          if exists(mask):
              mask = rearrange(mask, 'b ... -> b (...)')
              max_neg_value = -torch.finfo(sim.dtype).max
              mask = repeat(mask, 'b j -> (b h) () j', h=h)
              sim.masked_fill_(~mask, max_neg_value)
  
          # attention, what we cannot get enough of
          sim = sim.softmax(dim=-1)
  
          out = einsum('b i j, b j d -> b i d', sim, v)
          """
      ## new
      with sdp_kernel(**BACKEND_MAP[self.backend]):
          # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape)
          out = F.scaled_dot_product_attention(
              q, k, v, attn_mask=mask
          )  # scale is dim_head ** -0.5 per default
  
      del q, k, v
      out = rearrange(out, "b h n d -> b n (h d)", h=h)
  
      if additional_tokens is not None:
          # remove additional token
          out = out[:, n_tokens_to_mask:]
      return self.to_out(out)
  
imports: import logging
  import math
  from inspect import isfunction
  from typing import Any, Optional
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from packaging import version
  from torch import nn
  from torch.utils.checkpoint import checkpoint
label_desc: Use image node_id 2 for calling the CrossAttention obj's underlying member method code with example usage: obj.forward(x, context, mask, additional_tokens, n_times_crossframe_attn_in_) and returns: self