code: def forward(self, x, pos_emb, fmap_dims):
      b, n, _, h = *x.shape, self.heads
  
      to_q_kwargs = {"fmap_dims": fmap_dims} if self.use_ds_conv else {}
  
      x = self.norm(x)
  
      q = self.to_q(x, **to_q_kwargs)
  
      qkv = (q, *self.to_kv(x).chunk(2, dim=-1))
  
      q, k, v = map(
          lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), qkv
      )
  
      if self.use_rotary:
          # apply 2d rotary embeddings to queries and keys, excluding CLS tokens
  
          sin, cos = pos_emb
          dim_rotary = sin.shape[-1]
  
          (q_cls, q), (k_cls, k) = map(
              lambda t: (t[:, :1], t[:, 1:]), (q, k)
          )
  
          # handle the case where rotary dimension < head dimension
  
          (q, q_pass), (k, k_pass) = map(
              lambda t: (t[..., :dim_rotary], t[..., dim_rotary:]),
              (q, k),
          )
          q, k = map(
              lambda t: (t * cos) + (rotate_every_two(t) * sin), (q, k)
          )
          q, k = map(
              lambda t: torch.cat(t, dim=-1), ((q, q_pass), (k, k_pass))
          )
  
          # concat back the CLS tokens
  
          q = torch.cat((q_cls, q), dim=1)
          k = torch.cat((k_cls, k), dim=1)
  
      dots = einsum("b i d, b j d -> b i j", q, k) * self.scale
  
      attn = self.attend(dots)
      attn = self.dropout(attn)
  
      out = einsum("b i j, b j d -> b i d", attn, v)
      out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
      return self.to_out(out)
  
imports: from math import sqrt, pi, log
  import torch
  from torch import nn, einsum
  import torch.nn.functional
  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, pos_emb, fmap_dims) and returns: self