class_docstr: A single Transformer Encoder Block.
docstr: Args:
      input_embed_dim: The input embedding dimension
      num_heads: The number of attention heads
      ff_hidden_multiplier: The hidden dimension multiplier for the position-wise feed-forward layer
      ff_activation: The activation function for the position-wise feed-forward layer
      attn_dropout: The dropout probability for the attention layer
      keep_attn: Whether to keep the attention weights
      ff_dropout: The dropout probability for the position-wise feed-forward layer
      add_norm_dropout: The dropout probability for the residual connections
      transformer_head_dim: The dimension of the attention heads. If None, will default to input_embed_dim
docstr_tok: ["Args",":","input_embed_dim",":","The","input","embedding","dimension","num_heads",":","The","number","of","attention","heads","ff_hidden_multiplier",":","The","hidden","dimension","multiplier","for","the","position-wise","feed-forward","layer","ff_activation",":","The","activation","function","for","the","position-wise","feed-forward","layer","attn_dropout",":","The","dropout","probability","for","the","attention","layer","keep_attn",":","Whether","to","keep","the","attention","weights","ff_dropout",":","The","dropout","probability","for","the","position-wise","feed-forward","layer","add_norm_dropout",":","The","dropout","probability","for","the","residual","connections","transformer_head_dim",":","The","dimension","of","the","attention","heads",".","If","None",",","will","default","to","input_embed_dim"]
code: def __init__(
      self,
      input_embed_dim: int,
      num_heads: int = 8,
      ff_hidden_multiplier: int = 4,
      ff_activation: str = "GEGLU",
      attn_dropout: float = 0.1,
      keep_attn: bool = True,
      ff_dropout: float = 0.1,
      add_norm_dropout: float = 0.1,
      transformer_head_dim: Optional[int] = None,
  ):
      """
      Args:
          input_embed_dim: The input embedding dimension
          num_heads: The number of attention heads
          ff_hidden_multiplier: The hidden dimension multiplier for the position-wise feed-forward layer
          ff_activation: The activation function for the position-wise feed-forward layer
          attn_dropout: The dropout probability for the attention layer
          keep_attn: Whether to keep the attention weights
          ff_dropout: The dropout probability for the position-wise feed-forward layer
          add_norm_dropout: The dropout probability for the residual connections
          transformer_head_dim: The dimension of the attention heads. If None, will default to input_embed_dim
      """
      super().__init__()
      self.mha = MultiHeadedAttention(
          input_embed_dim,
          num_heads,
          head_dim=input_embed_dim
          if transformer_head_dim is None
          else transformer_head_dim,
          dropout=attn_dropout,
          keep_attn=keep_attn,
      )
  
      try:
          self.pos_wise_ff = GATED_UNITS[ff_activation](
              d_model=input_embed_dim,
              d_ff=input_embed_dim * ff_hidden_multiplier,
              dropout=ff_dropout,
          )
      except AttributeError:
          self.pos_wise_ff = PositionWiseFeedForward(
              d_model=input_embed_dim,
              d_ff=input_embed_dim * ff_hidden_multiplier,
              dropout=ff_dropout,
              activation=getattr(nn, self.hparams.ff_activation),
          )
      self.attn_add_norm = AddNorm(input_embed_dim, add_norm_dropout)
      self.ff_add_norm = AddNorm(input_embed_dim, add_norm_dropout)
  
imports: import math
  from typing import Optional
  import torch
  from einops import rearrange
  from torch import einsum, nn
  from pytorch_tabular.utils import _initialize_kaiming
  from .gated_units import GEGLU, PositionWiseFeedForward, ReGLU, SwiGLU
label_desc: Use image node_id 1 to create a new TransformerEncoderBlock object from inherited base classes: nn with example: obj = TransformerEncoderBlock(input_embed_dim, num_heads, ff_hidden_multiplier, ff_activation, attn_dropout, keep_attn, ff_dropout, add_norm_dropout, transformer_head_dim)