code: def forward(self, x):
      height, width, heads = *x.shape[-2:], self.heads
  
      x = self.norm(x)
  
      q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
  
      # split out heads
  
      q, k, v = map(
          lambda t: rearrange(t, "b (h d) ... -> b h (...) d", h=heads),
          (q, k, v),
      )
  
      # similarity
  
      dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
  
      # attention
  
      attn = self.attend(dots)
      attn = self.dropout(attn)
  
      # aggregate values
  
      out = torch.matmul(attn, v)
  
      # merge back heads
  
      out = rearrange(
          out, "b h (x y) d -> b (h d) x y", x=height, y=width
      )
      return self.to_out(out)
  
imports: from functools import partial
  import torch
  from torch import nn
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 2 for calling the ScalableSelfAttention obj's underlying member method code with example usage: obj.forward(x) and returns: self