code: def __init__(
      self,
      dim,
      fmap_size,
      heads=8,
      dim_key=32,
      dim_value=64,
      dropout=0.0,
      dim_out=None,
      downsample=False,
  ):
      super().__init__()
      inner_dim_key = dim_key * heads
      inner_dim_value = dim_value * heads
      dim_out = default(dim_out, dim)
  
      self.heads = heads
      self.scale = dim_key**-0.5
  
      self.to_q = nn.Sequential(
          nn.Conv2d(
              dim,
              inner_dim_key,
              1,
              stride=(2 if downsample else 1),
              bias=False,
          ),
          nn.BatchNorm2d(inner_dim_key),
      )
      self.to_k = nn.Sequential(
          nn.Conv2d(dim, inner_dim_key, 1, bias=False),
          nn.BatchNorm2d(inner_dim_key),
      )
      self.to_v = nn.Sequential(
          nn.Conv2d(dim, inner_dim_value, 1, bias=False),
          nn.BatchNorm2d(inner_dim_value),
      )
  
      self.attend = nn.Softmax(dim=-1)
      self.dropout = nn.Dropout(dropout)
  
      out_batch_norm = nn.BatchNorm2d(dim_out)
      nn.init.zeros_(out_batch_norm.weight)
  
      self.to_out = nn.Sequential(
          nn.GELU(),
          nn.Conv2d(inner_dim_value, dim_out, 1),
          out_batch_norm,
          nn.Dropout(dropout),
      )
  
      # positional bias
  
      self.pos_bias = nn.Embedding(fmap_size * fmap_size, heads)
  
      q_range = torch.arange(
          0, fmap_size, step=(2 if downsample else 1)
      )
      k_range = torch.arange(fmap_size)
  
      q_pos = torch.stack(
          torch.meshgrid(q_range, q_range, indexing="ij"), dim=-1
      )
      k_pos = torch.stack(
          torch.meshgrid(k_range, k_range, indexing="ij"), dim=-1
      )
  
      q_pos, k_pos = map(
          lambda t: rearrange(t, "i j c -> (i j) c"), (q_pos, k_pos)
      )
      rel_pos = (q_pos[:, None, ...] - k_pos[None, :, ...]).abs()
  
      x_rel, y_rel = rel_pos.unbind(dim=-1)
      pos_indices = (x_rel * fmap_size) + y_rel
  
      self.register_buffer("pos_indices", pos_indices)
  
imports: from math import ceil
  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 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(dim, fmap_size, heads, dim_key, dim_value, dropout, dim_out, downsample)