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","]",")"]
docstr: Completes the forward pass
  
  Arguments
  ---------
  x: torch.Tensor
      the data to be attended to
  
  Returns
  -------
  result: torch.Tensor
      The data, with attention applied
docstr_tok: ["Completes","the","forward","pass","Arguments","--","--","--","--","-","x",":","torch.Tensor","the","data","to","be","attended","to","Returns","--","--","--","-","result",":","torch.Tensor","The","data",",","with","attention","applied"]
code: def forward(self, x):
      """Completes the forward pass
  
      Arguments
      ---------
      x: torch.Tensor
          the data to be attended to
  
      Returns
      -------
      result: torch.Tensor
          The data, with attention applied
      """
      b, c, *spatial = x.shape
      x = x.reshape(b, c, -1)
      qkv = self.qkv(self.norm(x))
      h = self.attention(qkv)
      h = self.proj_out(h)
      return (x + h).reshape(b, c, *spatial)
  
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 2 for calling the AttentionBlock obj's underlying member method code with example usage: obj.forward(x) and returns: unknown