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: * `n_channels` is the number of channels in the input
  * `n_heads` is the number of heads in multi-head attention
  * `d_k` is the number of dimensions in each head
  * `n_groups` is the number of groups for [group normalization](../../normalization/group_norm/index.html)
code: def __init__(
      self,
      n_channels: int,
      n_heads: int = 1,
      d_k: int = None,
      n_groups: int = 32,
  ):
      """
      * `n_channels` is the number of channels in the input
      * `n_heads` is the number of heads in multi-head attention
      * `d_k` is the number of dimensions in each head
      * `n_groups` is the number of groups for [group normalization](../../normalization/group_norm/index.html)
      """
      super().__init__()
  
      # Default `d_k`
      if d_k is None:
          d_k = n_channels
      # Normalization layer
      self.norm = nn.GroupNorm(n_groups, n_channels)
      # Projections for query, key and values
      self.projection = nn.Linear(n_channels, n_heads * d_k * 3)
      # Linear layer for final transformation
      self.output = nn.Linear(n_heads * d_k, n_channels)
      # Scale for dot-product attention
      self.scale = d_k**-0.5
      #
      self.n_heads = n_heads
      self.d_k = d_k
  
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 1 to create a new AttentionBlock object from inherited base classes: Module with example: obj = AttentionBlock(n_channels, n_heads, d_k, n_groups)