class_docstr: ## Attention block
docstr: :param x: is the tensor of shape `[batch_size, channels, height, width]`
code: def forward(self, x: torch.Tensor):
      """
      :param x: is the tensor of shape `[batch_size, channels, height, width]`
      """
      # Normalize `x`
      x_norm = self.norm(x)
      # Get query, key and vector embeddings
      q = self.q(x_norm)
      k = self.k(x_norm)
      v = self.v(x_norm)
  
      # Reshape to query, key and vector embeedings from
      # `[batch_size, channels, height, width]` to
      # `[batch_size, channels, height * width]`
      b, c, h, w = q.shape
      q = q.view(b, c, h * w)
      k = k.view(b, c, h * w)
      v = v.view(b, c, h * w)
  
      # Compute $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)$
      attn = torch.einsum("bci,bcj->bij", q, k) * self.scale
      attn = F.softmax(attn, dim=2)
  
      # Compute $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$
      out = torch.einsum("bij,bcj->bci", attn, v)
  
      # Reshape back to `[batch_size, channels, height, width]`
      out = out.view(b, c, h, w)
      # Final $1 \times 1$ convolution layer
      out = self.proj_out(out)
  
      # Add residual connection
      return x + out
  
imports: from typing import List
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 2 for calling the AttnBlock obj's underlying member method code with example usage: obj.forward(x) and returns: unknown