code: def __init__(
      self,
      in_planes,
      out_planes,
      groups=8,
      kernel_size=56,
      stride=1,
      bias=False,
      width=False,
  ):
      if not ((in_planes % groups == 0) and (out_planes % groups == 0)):
          raise AssertionError
      super(AxialAttention, self).__init__()
      self.in_planes = in_planes
      self.out_planes = out_planes
      self.groups = groups
      self.group_planes = out_planes // groups
      self.kernel_size = kernel_size
      self.stride = stride
      self.bias = bias
      self.width = width
  
      # Multi-head self attention
      self.qkv_transform = qkv_transform(
          in_planes,
          out_planes * 2,
          kernel_size=1,
          stride=1,
          padding=0,
          bias=False,
      )
      self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
      self.bn_similarity = nn.BatchNorm2d(groups * 3)
      # self.bn_qk = nn.BatchNorm2d(groups)
      # self.bn_qr = nn.BatchNorm2d(groups)
      # self.bn_kr = nn.BatchNorm2d(groups)
      self.bn_output = nn.BatchNorm1d(out_planes * 2)
  
      # Position embedding
      self.relative = nn.Parameter(
          torch.randn(self.group_planes * 2, kernel_size * 2 - 1),
          requires_grad=True,
      )
      query_index = torch.arange(kernel_size).unsqueeze(0)
      key_index = torch.arange(kernel_size).unsqueeze(1)
      relative_index = key_index - query_index + kernel_size - 1
      self.register_buffer("flatten_index", relative_index.view(-1))
      if stride > 1:
          self.pooling = nn.AvgPool2d(stride, stride=stride)
  
      self.reset_parameters()
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new AxialAttention object from inherited base classes: nn with example: obj = AxialAttention(in_planes, out_planes, groups, kernel_size, stride, bias, width)