class_docstr: ### Attention block
  
  This is similar to [transformer multi-head attention](../../transformers/mha.html).
class_docstr_tok: ["#","#","#","Attention","block","This","is","similar","to","[","transformer","multi-head","attention","]","(","..","\/","..","\/transformers\/mha.html",")","."]
docstr: * `x` has shape `[batch_size, in_channels, height, width]`
  * `t` has shape `[batch_size, time_channels]`
code: def forward(self, x: torch.Tensor, t: Optional[torch.Tensor] = None):
      """
      * `x` has shape `[batch_size, in_channels, height, width]`
      * `t` has shape `[batch_size, time_channels]`
      """
      # `t` is not used, but it's kept in the arguments because for the attention layer function signature
      # to match with `ResidualBlock`.
      _ = t
      # Get shape
      batch_size, n_channels, height, width = x.shape
      # Change `x` to shape `[batch_size, seq, n_channels]`
      x = x.view(batch_size, n_channels, -1).permute(0, 2, 1)
      # Get query, key, and values (concatenated) and shape it to `[batch_size, seq, n_heads, 3 * d_k]`
      qkv = self.projection(x).view(
          batch_size, -1, self.n_heads, 3 * self.d_k
      )
      # Split query, key, and values. Each of them will have shape `[batch_size, seq, n_heads, d_k]`
      q, k, v = torch.chunk(qkv, 3, dim=-1)
      # Calculate scaled dot-product $\frac{Q K^\top}{\sqrt{d_k}}$
      attn = torch.einsum("bihd,bjhd->bijh", q, k) * self.scale
      # Softmax along the sequence dimension $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
      attn = attn.softmax(dim=2)
      # Multiply by values
      res = torch.einsum("bijh,bjhd->bihd", attn, v)
      # Reshape to `[batch_size, seq, n_heads * d_k]`
      res = res.view(batch_size, -1, self.n_heads * self.d_k)
      # Transform to `[batch_size, seq, n_channels]`
      res = self.output(res)
  
      # Add skip connection
      res += x
  
      # Change to shape `[batch_size, in_channels, height, width]`
      res = res.permute(0, 2, 1).view(
          batch_size, n_channels, height, width
      )
  
      #
      return res
  
imports: import math
  from typing import Optional, Tuple, Union, List
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the AttentionBlock obj's underlying member method code with example usage: obj.forward(x, t) and returns: res