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 forward(self, x, encoder_out=None):
      if not self.is_overwritten and (
          self.overwrite_qkv or self.overwrite_linear
      ):
          self.set_weights(self)
          self.is_overwritten = True
  
      b, c, *spatial = x.shape
      hid_states = self.norm(x).reshape([b, c, -1])
  
      qkv = self.qkv(hid_states)
      bs, width, length = qkv.shape
      assert width % (3 * self.n_heads) == 0
      ch = width // (3 * self.n_heads)
      q, k, v = qkv.reshape([bs * self.n_heads, ch * 3, length]).split(
          ch, axis=1
      )
  
      if encoder_out is not None:
          encoder_kv = self.encoder_kv(encoder_out)
          assert encoder_kv.shape[1] == self.n_heads * ch * 2
          ek, ev = encoder_kv.reshape(
              [bs * self.n_heads, ch * 2, -1]
          ).split(ch, axis=1)
          k = paddle.concat([ek, k], axis=-1)
          v = paddle.concat([ev, v], axis=-1)
  
      scale = 1 / math.sqrt(math.sqrt(ch))
      weight = paddle.einsum(
          "bct,bcs->bts", q * scale, k * scale
      )  # More stable with f16 than dividing afterwards
      weight = F.softmax(weight.astype("float32"), axis=-1).astype(
          weight.dtype
      )
  
      a = paddle.einsum("bts,bcs->bct", weight, v)
      h = a.reshape([bs, -1, length])
  
      h = self.proj(h)
      h = h.reshape([b, c, *spatial])
  
      result = x + h
  
      result = result / self.rescale_output_factor
  
      return result
  
imports: import math
  from inspect import isfunction
  import numpy
  import paddle
  import paddle.nn
  import paddle.nn.functional
label_desc: Use image node_id 3 for calling the AttentionBlock obj's underlying member method code with example usage: obj.forward(x, encoder_out) and returns: result