class_docstr: Single transformer block, can be used for both encoder and decoder
class_docstr_tok: ["Single","transformer","block",",","can","be","used","for","both","encoder","and","decoder"]
code: def __init__(
      self,
      seq_length: int,
      num_heads: int,
      hidden_dim: int,
      mlp_dim: int,
      drop_out_rate: float = 0.1,
      drop_path_rate: float = 0.1,
      norm_layer: Callable[..., torch.nn.Module] = partial(
          nn.LayerNorm, eps=1e-6
      ),
      is_decoder: bool = False,
  ) -> None:
      super().__init__()
      if is_decoder:
          self.register_buffer(
              "look_ahead_mask",
              torch.triu(
                  torch.ones(
                      seq_length, seq_length, dtype=torch.float32
                  ),
                  diagonal=1,
              ),
          )
      else:
          self.look_ahead_mask = None
      self.is_decoder = is_decoder
      self.num_heads = num_heads
  
      # --- BLOCK 1 ---
      self.ln_1a = norm_layer(hidden_dim)
      self.self_attention = WindowMultiHeadAttention(
          hidden_dim,
          num_heads,
          attn_drop=drop_out_rate,
          proj_drop=drop_out_rate,
      )
      self.ln_1b = norm_layer(hidden_dim)
      self.mlp1 = MLP(
          in_features=hidden_dim,
          mlp_dim=mlp_dim,
          dropout=drop_out_rate,
      )
  
      # --- BLOCK 2 ---
      self.ln_2a = norm_layer(hidden_dim)
      self.cross_attention = WindowMultiHeadAttention(
          hidden_dim,
          num_heads,
          attn_drop=drop_out_rate,
          proj_drop=drop_out_rate,
      )
      self.ln_2b = norm_layer(hidden_dim)
      self.mlp2 = MLP(
          in_features=hidden_dim, mlp_dim=mlp_dim, dropout=drop_out_rate
      )
  
      self.drop_path = StochasticDepth(drop_path_rate)
  
imports: from functools import partial
  from typing import Callable, Optional
  import torch
  import torch.nn
  from torch import Tensor
  from .layers_utils import MLP, StochasticDepth, WindowMultiHeadAttention
label_desc: Use image node_id 1 to create a new TransformerBlock object from inherited base classes: nn with example: obj = TransformerBlock(seq_length, num_heads, hidden_dim, mlp_dim, drop_out_rate, drop_path_rate, norm_layer, is_decoder)