def __init__(self, column_idx: Dict[(str, int)], cat_embed_input: Optional[List[Tuple[(str, int)]]]=None, cat_embed_dropout: float=0.1, use_cat_bias: bool=False, cat_embed_activation: Optional[str]=None, full_embed_dropout: bool=False, shared_embed: bool=False, add_shared_embed: bool=False, frac_shared_embed: float=0.25, continuous_cols: Optional[List[str]]=None, cont_norm_layer: str=None, cont_embed_dropout: float=0.1, use_cont_bias: bool=True, cont_embed_activation: Optional[str]=None, input_dim: int=32, use_qkv_bias: bool=False, n_heads: int=8, n_blocks: int=2, attn_dropout: float=0.1, ff_dropout: float=0.2, transformer_activation: str='gelu', mlp_hidden_dims: Optional[List[int]]=None, mlp_activation: str='relu', mlp_dropout: float=0.1, mlp_batchnorm: bool=False, mlp_batchnorm_last: bool=False, mlp_linear_first: bool=True):
    super(SAINT, self).__init__(column_idx=column_idx, cat_embed_input=cat_embed_input, cat_embed_dropout=cat_embed_dropout, use_cat_bias=use_cat_bias, cat_embed_activation=cat_embed_activation, full_embed_dropout=full_embed_dropout, shared_embed=shared_embed, add_shared_embed=add_shared_embed, frac_shared_embed=frac_shared_embed, continuous_cols=continuous_cols, cont_norm_layer=cont_norm_layer, embed_continuous=True, cont_embed_dropout=cont_embed_dropout, use_cont_bias=use_cont_bias, cont_embed_activation=cont_embed_activation, input_dim=input_dim)
    self.column_idx = column_idx
    self.cat_embed_input = cat_embed_input
    self.cat_embed_dropout = cat_embed_dropout
    self.full_embed_dropout = full_embed_dropout
    self.shared_embed = shared_embed
    self.add_shared_embed = add_shared_embed
    self.frac_shared_embed = frac_shared_embed
    self.continuous_cols = continuous_cols
    self.cont_embed_activation = cont_embed_activation
    self.cont_embed_dropout = cont_embed_dropout
    self.cont_norm_layer = cont_norm_layer
    self.input_dim = input_dim
    self.use_qkv_bias = use_qkv_bias
    self.n_heads = n_heads
    self.n_blocks = n_blocks
    self.attn_dropout = attn_dropout
    self.ff_dropout = ff_dropout
    self.transformer_activation = transformer_activation
    self.mlp_hidden_dims = mlp_hidden_dims
    self.mlp_activation = mlp_activation
    self.mlp_dropout = mlp_dropout
    self.mlp_batchnorm = mlp_batchnorm
    self.mlp_batchnorm_last = mlp_batchnorm_last
    self.mlp_linear_first = mlp_linear_first
    self.with_cls_token = ('cls_token' in column_idx)
    self.n_cat = (len(cat_embed_input) if (cat_embed_input is not None) else 0)
    self.n_cont = (len(continuous_cols) if (continuous_cols is not None) else 0)
    self.n_feats = (self.n_cat + self.n_cont)
    self.saint_blks = nn.Sequential()
    for i in range(n_blocks):
        self.saint_blks.add_module(('saint_block' + str(i)), SaintEncoder(input_dim, n_heads, use_qkv_bias, attn_dropout, ff_dropout, transformer_activation, self.n_feats))
    attn_output_dim = (self.input_dim if self.with_cls_token else (self.n_feats * self.input_dim))
    if (not mlp_hidden_dims):
        mlp_hidden_dims = [attn_output_dim, (attn_output_dim * 4), (attn_output_dim * 2)]
    else:
        mlp_hidden_dims = ([attn_output_dim] + mlp_hidden_dims)
    self.saint_mlp = MLP(mlp_hidden_dims, mlp_activation, mlp_dropout, mlp_batchnorm, mlp_batchnorm_last, mlp_linear_first)
    self.output_dim: int = mlp_hidden_dims[(- 1)]
def __init__(self, column_idx: Dict[(str, int)], cat_embed_input: Optional[List[Tuple[(str, int)]]]=None, cat_embed_dropout: float=0.1, use_cat_bias: bool=False, cat_embed_activation: Optional[str]=None, full_embed_dropout: bool=False, shared_embed: bool=False, add_shared_embed: bool=False, frac_shared_embed: float=0.25, continuous_cols: Optional[List[str]]=None, cont_norm_layer: str=None, cont_embed_dropout: float=0.1, use_cont_bias: bool=True, cont_embed_activation: Optional[str]=None, input_dim: int=32, use_qkv_bias: bool=False, n_heads: int=8, n_blocks: int=2, attn_dropout: float=0.1, ff_dropout: float=0.2, transformer_activation: str='gelu', mlp_hidden_dims: Optional[List[int]]=None, mlp_activation: str='relu', mlp_dropout: float=0.1, mlp_batchnorm: bool=False, mlp_batchnorm_last: bool=False, mlp_linear_first: bool=True):
    super(SAINT, self).__init__(column_idx=column_idx, cat_embed_input=cat_embed_input, cat_embed_dropout=cat_embed_dropout, use_cat_bias=use_cat_bias, cat_embed_activation=cat_embed_activation, full_embed_dropout=full_embed_dropout, shared_embed=shared_embed, add_shared_embed=add_shared_embed, frac_shared_embed=frac_shared_embed, continuous_cols=continuous_cols, cont_norm_layer=cont_norm_layer, embed_continuous=True, cont_embed_dropout=cont_embed_dropout, use_cont_bias=use_cont_bias, cont_embed_activation=cont_embed_activation, input_dim=input_dim)
    self.column_idx = column_idx
    self.cat_embed_input = cat_embed_input
    self.cat_embed_dropout = cat_embed_dropout
    self.full_embed_dropout = full_embed_dropout
    self.shared_embed = shared_embed
    self.add_shared_embed = add_shared_embed
    self.frac_shared_embed = frac_shared_embed
    self.continuous_cols = continuous_cols
    self.cont_embed_activation = cont_embed_activation
    self.cont_embed_dropout = cont_embed_dropout
    self.cont_norm_layer = cont_norm_layer
    self.input_dim = input_dim
    self.use_qkv_bias = use_qkv_bias
    self.n_heads = n_heads
    self.n_blocks = n_blocks
    self.attn_dropout = attn_dropout
    self.ff_dropout = ff_dropout
    self.transformer_activation = transformer_activation
    self.mlp_hidden_dims = mlp_hidden_dims
    self.mlp_activation = mlp_activation
    self.mlp_dropout = mlp_dropout
    self.mlp_batchnorm = mlp_batchnorm
    self.mlp_batchnorm_last = mlp_batchnorm_last
    self.mlp_linear_first = mlp_linear_first
    self.with_cls_token = ('cls_token' in column_idx)
    self.n_cat = (len(cat_embed_input) if (cat_embed_input is not None) else 0)
    self.n_cont = (len(continuous_cols) if (continuous_cols is not None) else 0)
    self.n_feats = (self.n_cat + self.n_cont)
    self.saint_blks = nn.Sequential()
    for i in range(n_blocks):
        self.saint_blks.add_module(('saint_block' + str(i)), SaintEncoder(input_dim, n_heads, use_qkv_bias, attn_dropout, ff_dropout, transformer_activation, self.n_feats))
    attn_output_dim = (self.input_dim if self.with_cls_token else (self.n_feats * self.input_dim))
    if (not mlp_hidden_dims):
        mlp_hidden_dims = [attn_output_dim, (attn_output_dim * 4), (attn_output_dim * 2)]
    else:
        mlp_hidden_dims = ([attn_output_dim] + mlp_hidden_dims)
    self.saint_mlp = MLP(mlp_hidden_dims, mlp_activation, mlp_dropout, mlp_batchnorm, mlp_batchnorm_last, mlp_linear_first)
    self.output_dim: int = mlp_hidden_dims[(- 1)]