code: def forward(self, x):
      if self.width:
          x = x.permute(0, 2, 1, 3)
      else:
          x = x.permute(0, 3, 1, 2)  # N, W, C, H
      N, W, C, H = x.shape
      x = x.contiguous().view(N * W, C, H)
  
      # Transformations
      qkv = self.bn_qkv(self.qkv_transform(x))
      q, k, v = torch.split(
          qkv.reshape(N * W, self.groups, self.group_planes * 2, H),
          [
              self.group_planes // 2,
              self.group_planes // 2,
              self.group_planes,
          ],
          dim=2,
      )
  
      # Calculate position embedding
      all_embeddings = torch.index_select(
          self.relative, 1, self.flatten_index
      ).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
      q_embedding, k_embedding, v_embedding = torch.split(
          all_embeddings,
          [
              self.group_planes // 2,
              self.group_planes // 2,
              self.group_planes,
          ],
          dim=0,
      )
      qr = torch.einsum("bgci,cij->bgij", q, q_embedding)
      kr = torch.einsum("bgci,cij->bgij", k, k_embedding).transpose(
          2, 3
      )
      qk = torch.einsum("bgci, bgcj->bgij", q, k)
      stacked_similarity = torch.cat([qk, qr, kr], dim=1)
      stacked_similarity = (
          self.bn_similarity(stacked_similarity)
          .view(N * W, 3, self.groups, H, H)
          .sum(dim=1)
      )
      # stacked_similarity = self.bn_qr(qr) + self.bn_kr(kr) + self.bn_qk(qk)
      # (N, groups, H, H, W)
      similarity = F.softmax(stacked_similarity, dim=3)
      sv = torch.einsum("bgij,bgcj->bgci", similarity, v)
      sve = torch.einsum("bgij,cij->bgci", similarity, v_embedding)
      stacked_output = torch.cat([sv, sve], dim=-1).view(
          N * W, self.out_planes * 2, H
      )
      output = (
          self.bn_output(stacked_output)
          .view(N, W, self.out_planes, 2, H)
          .sum(dim=-2)
      )
  
      if self.width:
          output = output.permute(0, 2, 1, 3)
      else:
          output = output.permute(0, 2, 3, 1)
  
      if self.stride > 1:
          output = self.pooling(output)
  
      return output
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the AxialAttention obj's underlying member method code with example usage: obj.forward(x) and returns: output