docstr: TabNet Will output a vector of size output_dim.
  
  Args:
      input_size: concatenated size of input feature encoder outputs
      size: Embedding feature dimension
      output_size: Output dimension for TabNet
      num_steps: Total number of steps.
      num_total_blocks: Total number of feature transformer blocks.
      num_shared_blocks: Number of shared feature transformer blocks.
      relaxation_factor: >1 will allow features to be used more than once.
      bn_momentum: Batch normalization, momentum.
      bn_epsilon: Batch normalization, epsilon.
      bn_virtual_bs: Virtual batch ize for ghost batch norm.
      entmax_mode: Entmax is a sparse family of probability mapping which generalizes softmax and sparsemax.
                   entmax_mode controls the sparsity.  One of {"sparsemax", "entmax15", "constant", "adaptive"}.
      entmax_alpha: Must be a number between 1.0 and 2.0.  If entmax_mode is "adaptive", entmax_alpha is used
                    as the initial value for the learnable parameter.
docstr_tok: ["TabNet","Will","output","a","vector","of","size","output_dim",".","Args",":","input_size",":","concatenated","size","of","input","feature","encoder","outputs","size",":","Embedding","feature","dimension","output_size",":","Output","dimension","for","TabNet","num_steps",":","Total","number","of","steps",".","num_total_blocks",":","Total","number","of","feature","transformer","blocks",".","num_shared_blocks",":","Number","of","shared","feature","transformer","blocks",".","relaxation_factor",":",">","1","will","allow","features","to","be","used","more","than","once",".","bn_momentum",":","Batch","normalization",",","momentum",".","bn_epsilon",":","Batch","normalization",",","epsilon",".","bn_virtual_bs",":","Virtual","batch","ize","for","ghost","batch","norm",".","entmax_mode",":","Entmax","is","a","sparse","family","of","probability","mapping","which","generalizes","softmax","and","sparsemax",".","entmax_mode","controls","the","sparsity",".","One","of","{","``","sparsemax","''",",","``","entmax15","''",",","``","constant","''",",","``","adaptive","''","}",".","entmax_alpha",":","Must","be","a","number","between","1.0","and","2.0",".","If","entmax_mode","is","``","adaptive","''",",","entmax_alpha","is","used","as","the","initial","value","for","the","learnable","parameter","."]
code: def __init__(
      self,
      input_size: int,
      size: int,
      output_size: int,
      num_steps: int = 1,
      num_total_blocks: int = 4,
      num_shared_blocks: int = 2,
      relaxation_factor: float = 1.5,
      bn_momentum: float = 0.3,
      bn_epsilon: float = 1e-3,
      bn_virtual_bs: Optional[int] = None,
      sparsity: float = 1e-5,
      entmax_mode: str = "sparsemax",
      entmax_alpha: float = 1.5,
  ):
      """TabNet Will output a vector of size output_dim.
  
      Args:
          input_size: concatenated size of input feature encoder outputs
          size: Embedding feature dimension
          output_size: Output dimension for TabNet
          num_steps: Total number of steps.
          num_total_blocks: Total number of feature transformer blocks.
          num_shared_blocks: Number of shared feature transformer blocks.
          relaxation_factor: >1 will allow features to be used more than once.
          bn_momentum: Batch normalization, momentum.
          bn_epsilon: Batch normalization, epsilon.
          bn_virtual_bs: Virtual batch ize for ghost batch norm.
          entmax_mode: Entmax is a sparse family of probability mapping which generalizes softmax and sparsemax.
                       entmax_mode controls the sparsity.  One of {"sparsemax", "entmax15", "constant", "adaptive"}.
          entmax_alpha: Must be a number between 1.0 and 2.0.  If entmax_mode is "adaptive", entmax_alpha is used
                        as the initial value for the learnable parameter.
      """
      super().__init__()
      self.input_size = input_size
      self.size = size
      self.output_size = output_size
      self.num_steps = num_steps
      self.bn_virtual_bs = bn_virtual_bs
      self.relaxation_factor = relaxation_factor
      self.sparsity = torch.tensor(sparsity)
      self.batch_norm = nn.BatchNorm1d(
          input_size, momentum=bn_momentum, eps=bn_epsilon
      )
  
      kargs = {
          "num_total_blocks": num_total_blocks,
          "num_shared_blocks": num_shared_blocks,
          "bn_momentum": bn_momentum,
          "bn_epsilon": bn_epsilon,
          "bn_virtual_bs": bn_virtual_bs,
      }
  
      # first feature transformer block is built first
      # to get the shared blocks
      self.feature_transforms = nn.ModuleList(
          [FeatureTransformer(input_size, size + output_size, **kargs)]
      )
      self.attentive_transforms = nn.ModuleList([None])
      for i in range(num_steps):
          self.feature_transforms.append(
              FeatureTransformer(
                  input_size,
                  size + output_size,
                  **kargs,
                  shared_fc_layers=self.feature_transforms[
                      0
                  ].shared_fc_layers,
              )
          )
          # attentive transformers are initialized in build
          # because their outputs size depends on the number
          # of features that we determine by looking at the
          # last dimension of the input tensor
          self.attentive_transforms.append(
              AttentiveTransformer(
                  size,
                  input_size,
                  bn_momentum,
                  bn_epsilon,
                  bn_virtual_bs,
                  entmax_mode,
                  entmax_alpha,
              )
          )
      self.final_projection = nn.Linear(output_size, output_size)
  
      # Register tensors to be used in forward pass. This is needed in order to move these tensors
      # to the correct device (GPU/CPU) during the forward pass.
      self.register_buffer("out_accumulator", torch.zeros(output_size))
      self.register_buffer("aggregated_mask", torch.zeros(input_size))
      self.register_buffer("prior_scales", torch.ones(input_size))
  
imports: from typing import List, Optional, Tuple
  import torch
  import torch.nn
  from ludwig.modules.normalization_modules import GhostBatchNormalization
  from ludwig.utils.entmax import Entmax15, EntmaxBisect, Sparsemax
  from ludwig.utils.torch_utils import LudwigModule
label_desc: Use image node_id 1 to create a new TabNet object from inherited base classes: LudwigModule with example: obj = TabNet(input_size, size, output_size, num_steps, num_total_blocks, num_shared_blocks, relaxation_factor, bn_momentum, bn_epsilon, bn_virtual_bs, sparsity, entmax_mode, entmax_alpha)