class_docstr: An attention block that allows spatial positions to attend to each other.
  Originally ported from here, but adapted to the N-d case.
  https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
  
  
  Arguments
  ---------
  channels: int
      the number of channels
  num_heads: int
      the number of attention heads
  num_head_channels: int
      the number of channels in each attention head
  norm_num_groups: int
      the number of groups used for group normalization
  use_fixup_init: bool
      whether to use FixUp initialization
  
  Example
  -------
  >>> attn = AttentionBlock(
  ...     channels=8,
  ...     num_heads=4,
  ...     num_head_channels=4,
  ...     norm_num_groups=2
  ... )
  >>> x = torch.randn(4, 8, 16, 16)
  >>> out = attn(x)
  >>> out.shape
  torch.Size([4, 8, 16, 16])
class_docstr_tok: ["An","attention","block","that","allows","spatial","positions","to","attend","to","each","other",".","Originally","ported","from","here",",","but","adapted","to","the","N-d","case",".","https",":","\/\/github.com\/hojonathanho\/diffusion\/blob\/1e0dceb3b3495bbe19116a5e1b3596cd0706c543\/diffusion_tf\/models\/unet.py","#","L66",".","Arguments","--","--","--","--","-","channels",":","int","the","number","of","channels","num_heads",":","int","the","number","of","attention","heads","num_head_channels",":","int","the","number","of","channels","in","each","attention","head","norm_num_groups",":","int","the","number","of","groups","used","for","group","normalization","use_fixup_init",":","bool","whether","to","use","FixUp","initialization","Example","--","--","--","-",">",">",">","attn","=","AttentionBlock","(","...","channels=8",",","...","num_heads=4",",","...","num_head_channels=4",",","...","norm_num_groups=2","...",")",">",">",">","x","=","torch.randn","(","4",",","8",",","16",",","16",")",">",">",">","out","=","attn","(","x",")",">",">",">","out.shape","torch.Size","(","[","4",",","8",",","16",",","16","]",")"]
code: def __init__(
      self,
      channels,
      num_heads=1,
      num_head_channels=-1,
      norm_num_groups=32,
      use_fixup_init=True,
  ):
      super().__init__()
      self.channels = channels
      if num_head_channels == -1:
          self.num_heads = num_heads
      else:
          assert (
              channels % num_head_channels == 0
          ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
          self.num_heads = channels // num_head_channels
      self.norm = nn.GroupNorm(norm_num_groups, channels)
      self.qkv = conv_nd(1, channels, channels * 3, 1)
      self.attention = QKVAttention(self.num_heads)
  
      self.proj_out = fixup(
          conv_nd(1, channels, channels, 1), use_fixup_init
      )
  
imports: from abc import abstractmethod
  from speechbrain.utils.data_utils import pad_divisible
  from .autoencoders import NormalizingAutoencoder
  import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new AttentionBlock object from inherited base classes: nn with example: obj = AttentionBlock(channels, num_heads, num_head_channels, norm_num_groups, use_fixup_init)