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.
  Uses three q, k, v linear layers to compute attention
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",".","Uses","three","q",",","k",",","v","linear","layers","to","compute","attention"]
code: def forward(self, hidden_states):
      residual = hidden_states
      batch, channel, height, width = hidden_states.shape
  
      # norm
      hidden_states = self.group_norm(hidden_states)
  
      hidden_states = hidden_states.reshape(
          [batch, channel, height * width]
      ).transpose([0, 2, 1])
  
      # proj to q, k, v
      query_proj = self.query(hidden_states)
      key_proj = self.key(hidden_states)
      value_proj = self.value(hidden_states)
  
      # transpose
      query_states = self.transpose_for_scores(query_proj)
      key_states = self.transpose_for_scores(key_proj)
      value_states = self.transpose_for_scores(value_proj)
  
      # get scores
      scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads))
      attention_scores = paddle.matmul(
          query_states * scale, key_states * scale, transpose_y=True
      )
      attention_probs = F.softmax(
          attention_scores.astype("float32"), axis=-1
      ).astype(attention_scores.dtype)
  
      # compute attention output
      context_states = paddle.matmul(attention_probs, value_states)
  
      context_states = context_states.transpose([0, 2, 1, 3])
      new_context_states_shape = context_states.shape[:-2] + [
          self.channels,
      ]
      context_states = context_states.reshape(new_context_states_shape)
  
      # compute next hidden_states
      hidden_states = self.proj_attn(context_states)
      hidden_states = hidden_states.transpose([0, 2, 1]).reshape(
          [batch, channel, height, width]
      )
  
      # res connect and rescale
      hidden_states = (
          hidden_states + residual
      ) / self.rescale_output_factor
      return hidden_states
  
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 AttentionBlockNew obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: hidden_states