def __init__(
        self,
        input_sizes: dict[str, int],
        hidden_size: int,
        input_embedding_flags: Optional[dict[str, bool]] = None,
        dropout: float = 0.1,
        context_size: int = None,
        single_variable_grns: Optional[dict[str, _GatedResidualNetwork]] = None,
        prescalers: Optional[dict[str, nn.Linear]] = None,
        layer_norm: nn.Module = nn.LayerNorm,
    ):
        """
        Calculate weights for ``num_inputs`` variables  which are each of size ``input_size``
        """
        super().__init__()

        input_embedding_flags = (
            input_embedding_flags if input_embedding_flags is not None else {}
        )
        single_variable_grns = (
            single_variable_grns if single_variable_grns is not None else {}
        )
        prescalers = prescalers if prescalers is not None else {}

        self.hidden_size = hidden_size
        self.input_sizes = input_sizes
        self.input_embedding_flags = input_embedding_flags
        self.dropout = dropout
        self.context_size = context_size

        if self.num_inputs > 1:
            if self.context_size is not None:
                self.flattened_grn = _GatedResidualNetwork(
                    self.input_size_total,
                    min(self.hidden_size, self.num_inputs),
                    self.num_inputs,
                    self.dropout,
                    self.context_size,
                    residual=False,
                )
            else:
                self.flattened_grn = _GatedResidualNetwork(
                    self.input_size_total,
                    min(self.hidden_size, self.num_inputs),
                    self.num_inputs,
                    self.dropout,
                    residual=False,
                )

        self.single_variable_grns = nn.ModuleDict()
        self.prescalers = nn.ModuleDict()
        for name, input_size in self.input_sizes.items():
            if name in single_variable_grns:
                self.single_variable_grns[name] = single_variable_grns[name]
            elif self.input_embedding_flags.get(name, False):
                self.single_variable_grns[name] = _ResampleNorm(
                    input_size,
                    self.hidden_size,
                    norm=layer_norm,
                )
            else:
                self.single_variable_grns[name] = _GatedResidualNetwork(
                    input_size,
                    min(input_size, self.hidden_size),
                    output_size=self.hidden_size,
                    dropout=self.dropout,
                )
            if name in prescalers:  # reals need to be first scaled up
                self.prescalers[name] = prescalers[name]
            elif not self.input_embedding_flags.get(name, False):
                self.prescalers[name] = nn.Linear(1, input_size)

        self.softmax = nn.Softmax(dim=-1)

        Calculate weights for ``num_inputs`` variables  which are each of size ``input_size``
        
darts/models/forecasting/tft_submodels.py