class_docstr: Transformer with extra features.
  
  This module is the backbone of `FTTransformer`.
class_docstr_tok: ["Transformer","with","extra","features",".","This","module","is","the","backbone","of","`","FTTransformer","`","."]
docstr: Parameters
  ----------
  d_token
      The size of one token for `_CategoricalFeatureTokenizer`.
  n_blocks
      Number of the `FT_Transformer` blocks, which should be non-negative.
  attention_n_heads
      Number of attention heads in each `FT_Transformer` block, which should be positive.
  attention_dropout
      Dropout ratio for the Multi Headed Attention module.
  attention_initialization
      Weights initialization scheme for Multi Headed Attention module.
  attention_normalization
      Normalization policy for attention layers. "layer_norm" is a good default.
  ffn_d_hidden
      Number of the hidden nodes of the linear layers in the Feed-Forward Network module.
  ffn_dropout
      Dropout ratio of the hidden nodes of the linear layers in the Feed-Forward Network module.
  ffn_activation
      Activation function type for the Feed-Forward Network module.
  ffn_normalization
      Normalization scheme of the Feed-Forward Network module.
  residual_dropout
      Dropout ratio for the linear layers in FT_Transformer block.
  prenormalization, first_prenormalization
      Prenormalization to stabilize the training.
  n_tokens
      Number of tokens of the input sequence.
  kv_compression_ratio
      The compression ration to reduce the input sequence length.
  kv_compression_sharing
      If `true` the projections will share weights.
  head_activation
      Activation function type of the MLP layer.
  head_normalization
      Normalization scheme of the MLP layer.
  d_out
      Output dimension.
  projection
      Whether to use a project head.
  additive_attention
      If 'true' the transformer will use additive attention with linear complexity to sequence length.
  share_qv_weights
      if 'true', then value and query transformation parameters are shared in additive attention.
code: def __init__(
      self,
      *,
      d_token: int,
      n_blocks: int,
      attention_n_heads: int,
      attention_dropout: float,
      attention_initialization: str,
      attention_normalization: str,
      ffn_d_hidden: int,
      ffn_dropout: float,
      ffn_activation: str,
      ffn_normalization: str,
      residual_dropout: float,
      prenormalization: bool,
      first_prenormalization: bool,
      last_layer_query_idx: Union[None, List[int], slice],
      n_tokens: Optional[int],
      kv_compression_ratio: Optional[float],
      kv_compression_sharing: Optional[str],
      head_activation: ModuleType,
      head_normalization: ModuleType,
      d_out: int,
      projection: Optional[bool] = False,
      additive_attention: Optional[bool] = False,
      share_qv_weights: Optional[bool] = False,
  ) -> None:
      """
      Parameters
      ----------
      d_token
          The size of one token for `_CategoricalFeatureTokenizer`.
      n_blocks
          Number of the `FT_Transformer` blocks, which should be non-negative.
      attention_n_heads
          Number of attention heads in each `FT_Transformer` block, which should be positive.
      attention_dropout
          Dropout ratio for the Multi Headed Attention module.
      attention_initialization
          Weights initialization scheme for Multi Headed Attention module.
      attention_normalization
          Normalization policy for attention layers. "layer_norm" is a good default.
      ffn_d_hidden
          Number of the hidden nodes of the linear layers in the Feed-Forward Network module.
      ffn_dropout
          Dropout ratio of the hidden nodes of the linear layers in the Feed-Forward Network module.
      ffn_activation
          Activation function type for the Feed-Forward Network module.
      ffn_normalization
          Normalization scheme of the Feed-Forward Network module.
      residual_dropout
          Dropout ratio for the linear layers in FT_Transformer block.
      prenormalization, first_prenormalization
          Prenormalization to stabilize the training.
      n_tokens
          Number of tokens of the input sequence.
      kv_compression_ratio
          The compression ration to reduce the input sequence length.
      kv_compression_sharing
          If `true` the projections will share weights.
      head_activation
          Activation function type of the MLP layer.
      head_normalization
          Normalization scheme of the MLP layer.
      d_out
          Output dimension.
      projection
          Whether to use a project head.
      additive_attention
          If 'true' the transformer will use additive attention with linear complexity to sequence length.
      share_qv_weights
          if 'true', then value and query transformation parameters are shared in additive attention.
      """
      super().__init__()
      if isinstance(last_layer_query_idx, int):
          raise ValueError(
              "last_layer_query_idx must be None, list[int] or slice. "
              f"Do you mean last_layer_query_idx=[{last_layer_query_idx}] ?"
          )
      if not prenormalization:
          assert (
              not first_prenormalization
          ), "If `prenormalization` is False, then `first_prenormalization` must be False"
      assert _all_or_none(
          [n_tokens, kv_compression_ratio, kv_compression_sharing]
      ), (
          "If any of the following arguments is (not) None, then all of them must (not) be None: "
          "n_tokens, kv_compression_ratio, kv_compression_sharing"
      )
      assert (
          additive_attention or not share_qv_weights
      ), "If `share_qv_weights` is True, then `additive_attention` must be True"
      assert kv_compression_sharing in [
          None,
          "headwise",
          "key-value",
          "layerwise",
      ]
      if not prenormalization:
          if self.WARNINGS["prenormalization"]:
              warnings.warn(
                  "prenormalization is set to False. Are you sure about this? "
                  "The training can become less stable. "
                  "You can turn off this warning by tweaking the "
                  "rtdl.Transformer.WARNINGS dictionary.",
                  UserWarning,
              )
          assert (
              not first_prenormalization
          ), "If prenormalization is False, then first_prenormalization is ignored and must be set to False"
      if (
          prenormalization
          and first_prenormalization
          and self.WARNINGS["first_prenormalization"]
      ):
          warnings.warn(
              "first_prenormalization is set to True. Are you sure about this? "
              "For example, the vanilla FTTransformer with "
              "first_prenormalization=True performs SIGNIFICANTLY worse. "
              "You can turn off this warning by tweaking the "
              "rtdl.Transformer.WARNINGS dictionary.",
              UserWarning,
          )
  
      def make_kv_compression():
          assert (
              n_tokens and kv_compression_ratio
          ), _INTERNAL_ERROR_MESSAGE  # for mypy
          # https://github.com/pytorch/fairseq/blob/1bba712622b8ae4efb3eb793a8a40da386fe11d0/examples/linformer/linformer_src/modules/multihead_linear_attention.py#L83
          return nn.Linear(
              n_tokens, int(n_tokens * kv_compression_ratio), bias=False
          )
  
      self.shared_kv_compression = (
          make_kv_compression()
          if kv_compression_ratio
          and kv_compression_sharing == "layerwise"
          else None
      )
  
      self.prenormalization = prenormalization
      self.last_layer_query_idx = last_layer_query_idx
  
      self.blocks = nn.ModuleList([])
      for layer_idx in range(n_blocks):
          layer = nn.ModuleDict(
              {
                  "attention": AdditiveAttention(
                      d_token=d_token,
                      n_heads=attention_n_heads,
                      dropout=attention_dropout,
                      bias=True,
                      share_qv_weights=share_qv_weights,
                      initialization=attention_initialization,
                  )
                  if additive_attention
                  else MultiheadAttention(
                      d_token=d_token,
                      n_heads=attention_n_heads,
                      dropout=attention_dropout,
                      bias=True,
                      initialization=attention_initialization,
                  ),
                  "ffn": Custom_Transformer.FFN(
                      d_token=d_token,
                      d_hidden=ffn_d_hidden,
                      bias_first=True,
                      bias_second=True,
                      dropout=ffn_dropout,
                      activation=ffn_activation,
                  ),
                  "attention_residual_dropout": nn.Dropout(
                      residual_dropout
                  ),
                  "ffn_residual_dropout": nn.Dropout(residual_dropout),
                  "output": nn.Identity(),  # for hooks-based introspection
              }
          )
          if (
              layer_idx
              or not prenormalization
              or first_prenormalization
          ):
              layer["attention_normalization"] = _make_nn_module(
                  attention_normalization, d_token
              )
          layer["ffn_normalization"] = _make_nn_module(
              ffn_normalization, d_token
          )
          if (
              kv_compression_ratio
              and self.shared_kv_compression is None
          ):
              layer["key_compression"] = make_kv_compression()
              if kv_compression_sharing == "headwise":
                  layer["value_compression"] = make_kv_compression()
              else:
                  assert (
                      kv_compression_sharing == "key-value"
                  ), _INTERNAL_ERROR_MESSAGE
          self.blocks.append(layer)
  
      self.head = (
          Custom_Transformer.Head(
              d_in=d_token,
              d_out=d_out,
              bias=True,
              activation=head_activation,  # type: ignore
              normalization=head_normalization
              if prenormalization
              else "Identity",
          )
          if projection
          else nn.Identity()
      )
  
imports: import enum
  import math
  import warnings
  from typing import Callable, Dict, List, Optional, Tuple, Union, cast
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from .utils import init_weights
label_desc: Use image node_id 1 to create a new Custom_Transformer object from inherited base classes: nn with example: obj = Custom_Transformer()