code: def __init__(self, dim, dim_head=32, dropout=0.0, window_size=7):
      super().__init__()
      assert (
          dim % dim_head
      ) == 0, "dimension should be divisible by dimension per head"
  
      self.heads = dim // dim_head
      self.scale = dim_head**-0.5
  
      self.norm = nn.LayerNorm(dim)
      self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
  
      self.attend = nn.Sequential(
          nn.Softmax(dim=-1), nn.Dropout(dropout)
      )
  
      self.to_out = nn.Sequential(
          nn.Linear(dim, dim, bias=False), nn.Dropout(dropout)
      )
  
      # relative positional bias
  
      self.rel_pos_bias = nn.Embedding(
          (2 * window_size - 1) ** 2, self.heads
      )
  
      pos = torch.arange(window_size)
      grid = torch.stack(torch.meshgrid(pos, pos, indexing="ij"))
      grid = rearrange(grid, "c i j -> (i j) c")
      rel_pos = rearrange(grid, "i ... -> i 1 ...") - rearrange(
          grid, "j ... -> 1 j ..."
      )
      rel_pos += window_size - 1
      rel_pos_indices = (
          rel_pos * torch.tensor([2 * window_size - 1, 1])
      ).sum(dim=-1)
  
      self.register_buffer(
          "rel_pos_indices", rel_pos_indices, persistent=False
      )
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(dim, dim_head, dropout, window_size)