class_docstr: A modification of BasicTransformerBlock which supports pre-LayerNorm and post-LayerNorm configurations.
  
  Parameters:
      dim (`int`): The number of channels in the input and output.
      num_attention_heads (`int`): The number of heads to use for multi-head attention.
      attention_head_dim (`int`): The number of channels in each head.
      dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
      cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
      activation_fn (`str`, *optional*, defaults to `"geglu"`):
          Activation function to be used in feed-forward.
      num_embeds_ada_norm (:obj: `int`, *optional*):
          The number of diffusion steps used during training. See `Transformer2DModel`.
      attention_bias (:obj: `bool`, *optional*, defaults to `False`):
          Configure if the attentions should contain a bias parameter.
      only_cross_attention (`bool`, *optional*):
          Whether to use only cross-attention layers. In this case two cross attention layers are used.
      double_self_attention (`bool`, *optional*):
          Whether to use two self-attention layers. In this case no cross attention layers are used.
      upcast_attention (`bool`, *optional*):
          Whether to upcast the query and key to float32 when performing the attention calculation.
      norm_elementwise_affine (`bool`, *optional*):
          Whether to use learnable per-element affine parameters during layer normalization.
      norm_type (`str`, defaults to `"layer_norm"`):
          The layer norm implementation to use.
      pre_layer_norm (`bool`, *optional*):
          Whether to perform layer normalization before the attention and feedforward operations ("pre-LayerNorm"),
          as opposed to after ("post-LayerNorm"). Note that `BasicTransformerBlock` uses pre-LayerNorm, e.g.
          `pre_layer_norm = True`.
      final_dropout (`bool`, *optional*):
          Whether to use a final Dropout layer after the feedforward network.
class_docstr_tok: ["A","modification","of","BasicTransformerBlock","which","supports","pre-LayerNorm","and","post-LayerNorm","configurations",".","Parameters",":","dim","(","`","int","`",")",":","The","number","of","channels","in","the","input","and","output",".","num_attention_heads","(","`","int","`",")",":","The","number","of","heads","to","use","for","multi-head","attention",".","attention_head_dim","(","`","int","`",")",":","The","number","of","channels","in","each","head",".","dropout","(","`","float","`",",","*","optional","*",",","defaults","to","0.0",")",":","The","dropout","probability","to","use",".","cross_attention_dim","(","`","int","`",",","*","optional","*",")",":","The","size","of","the","encoder_hidden_states","vector","for","cross","attention",".","activation_fn","(","`","str","`",",","*","optional","*",",","defaults","to","`","``","geglu","''","`",")",":","Activation","function","to","be","used","in","feed-forward",".","num_embeds_ada_norm","(",":","obj",":","`","int","`",",","*","optional","*",")",":","The","number","of","diffusion","steps","used","during","training",".","See","`","Transformer2DModel","`",".","attention_bias","(",":","obj",":","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Configure","if","the","attentions","should","contain","a","bias","parameter",".","only_cross_attention","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","only","cross-attention","layers",".","In","this","case","two","cross","attention","layers","are","used",".","double_self_attention","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","two","self-attention","layers",".","In","this","case","no","cross","attention","layers","are","used",".","upcast_attention","(","`","bool","`",",","*","optional","*",")",":","Whether","to","upcast","the","query","and","key","to","float32","when","performing","the","attention","calculation",".","norm_elementwise_affine","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","learnable","per-element","affine","parameters","during","layer","normalization",".","norm_type","(","`","str","`",",","defaults","to","`","``","layer_norm","''","`",")",":","The","layer","norm","implementation","to","use",".","pre_layer_norm","(","`","bool","`",",","*","optional","*",")",":","Whether","to","perform","layer","normalization","before","the","attention","and","feedforward","operations","(","``","pre-LayerNorm","''",")",",","as","opposed","to","after","(","``","post-LayerNorm","''",")",".","Note","that","`","BasicTransformerBlock","`","uses","pre-LayerNorm",",","e.g",".","`","pre_layer_norm","=","True","`",".","final_dropout","(","`","bool","`",",","*","optional","*",")",":","Whether","to","use","a","final","Dropout","layer","after","the","feedforward","network","."]
code: def forward(
      self,
      hidden_states,
      attention_mask=None,
      encoder_hidden_states=None,
      encoder_attention_mask=None,
      timestep=None,
      cross_attention_kwargs=None,
      class_labels=None,
  ):
      # Pre-LayerNorm
      if self.pre_layer_norm:
          if self.use_ada_layer_norm:
              norm_hidden_states = self.norm1(hidden_states, timestep)
          else:
              norm_hidden_states = self.norm1(hidden_states)
      else:
          norm_hidden_states = hidden_states
  
      # 1. Self-Attention
      cross_attention_kwargs = (
          cross_attention_kwargs
          if cross_attention_kwargs is not None
          else {}
      )
      attn_output = self.attn1(
          norm_hidden_states,
          encoder_hidden_states=encoder_hidden_states
          if self.only_cross_attention
          else None,
          attention_mask=attention_mask,
          **cross_attention_kwargs,
      )
  
      # Post-LayerNorm
      if not self.pre_layer_norm:
          if self.use_ada_layer_norm:
              attn_output = self.norm1(attn_output, timestep)
          else:
              attn_output = self.norm1(attn_output)
  
      hidden_states = attn_output + hidden_states
  
      if self.attn2 is not None:
          # Pre-LayerNorm
          if self.pre_layer_norm:
              norm_hidden_states = (
                  self.norm2(hidden_states, timestep)
                  if self.use_ada_layer_norm
                  else self.norm2(hidden_states)
              )
          else:
              norm_hidden_states = hidden_states
          # TODO (Birch-San): Here we should prepare the encoder_attention mask correctly
          # prepare attention mask here
  
          # 2. Cross-Attention
          attn_output = self.attn2(
              norm_hidden_states,
              encoder_hidden_states=encoder_hidden_states,
              attention_mask=encoder_attention_mask,
              **cross_attention_kwargs,
          )
  
          # Post-LayerNorm
          if not self.pre_layer_norm:
              attn_output = (
                  self.norm2(attn_output, timestep)
                  if self.use_ada_layer_norm
                  else self.norm2(attn_output)
              )
  
          hidden_states = attn_output + hidden_states
  
      # 3. Feed-forward
      # Pre-LayerNorm
      if self.pre_layer_norm:
          norm_hidden_states = self.norm3(hidden_states)
      else:
          norm_hidden_states = hidden_states
  
      ff_output = self.ff(norm_hidden_states)
  
      # Post-LayerNorm
      if not self.pre_layer_norm:
          ff_output = self.norm3(ff_output)
  
      hidden_states = ff_output + hidden_states
  
      return hidden_states
  
imports: import math
  from typing import Optional, Union
  import torch
  from torch import nn
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  from ...models.attention import AdaLayerNorm, FeedForward
  from ...models.attention_processor import Attention
  from ...models.embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed
  from ...models.transformer_2d import Transformer2DModelOutput
  from ...utils import logging
label_desc: Use image node_id 2 for calling the UTransformerBlock obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, encoder_hidden_states, encoder_attention_mask, timestep, cross_attention_kwargs, class_labels) and returns: hidden_states