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.
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","."]
code: def __init__(
      self,
      channels: int,
      num_heads: int = 1,
      num_head_channels: int = -1,
      use_checkpoint: bool = False,
      use_new_attention_order: bool = False,
  ):
      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.use_checkpoint = use_checkpoint
      self.norm = normalization(channels)
      self.qkv = conv_nd(1, channels, channels * 3, 1)
      if use_new_attention_order:
          # split qkv before split heads
          self.attention = QKVAttention(self.num_heads)
      else:
          # split heads before split qkv
          self.attention = QKVAttentionLegacy(self.num_heads)
  
      self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
  
imports: import logging
  import math
  from abc import abstractmethod
  from typing import Iterable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from torch.utils.checkpoint import checkpoint
  from ...modules.attention import SpatialTransformer
  from ...modules.diffusionmodules.util import avg_pool_nd, conv_nd, linear, normalization, timestep_embedding, zero_module
  from ...modules.video_attention import SpatialVideoTransformer
  from ...util import exists
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, use_checkpoint, use_new_attention_order)