docstr: PyTorch module implementing the TFT architecture from `this paper <https://arxiv.org/pdf/1912.09363.pdf>`_
  The implementation is built upon `pytorch-forecasting's TemporalFusionTransformer
  <https://pytorch-forecasting.readthedocs.io/en/latest/models.html>`_.
  
  Parameters
  ----------
  output_dim : Tuple[int, int]
      shape of output given by (n_targets, loss_size). (loss_size corresponds to nr_params in other models).
  variables_meta : Dict[str, Dict[str, List[str]]]
      dict containing variable encoder, decoder variable names for mapping tensors in `_TFTModule.forward()`
  num_static_components
      the number of static components (not variables) of the input target series. This is either equal to the
      number of target components or 1.
  hidden_size : int
      hidden state size of the TFT. It is the main hyper-parameter and common across the internal TFT
      architecture.
  lstm_layers : int
      number of layers for the Long Short Term Memory (LSTM) Encoder and Decoder (1 is a good default).
  num_attention_heads : int
      number of attention heads (4 is a good default)
  full_attention : bool
      If `True`, applies multi-head attention query on past (encoder) and future (decoder) parts. Otherwise,
      only queries on future part. Defaults to `False`.
  feed_forward
      Set the feedforward network block. default `GatedResidualNetwork` or one of the  glu variant.
      Defaults to `GatedResidualNetwork`.
  hidden_continuous_size : int
      default for hidden size for processing continuous variables.
  categorical_embedding_sizes : dict
      A dictionary containing embedding sizes for categorical static covariates. The keys are the column names
      of the categorical static covariates. The values are tuples of integers with
      `(number of unique categories, embedding size)`. For example `{"some_column": (64, 8)}`.
      Note that `TorchForecastingModels` can only handle numeric data. Consider transforming/encoding your data
      with `darts.dataprocessing.transformers.static_covariates_transformer.StaticCovariatesTransformer`.
  dropout : float
      Fraction of neurons affected by Dropout.
  add_relative_index : bool
      Whether to add positional values to future covariates. Defaults to `False`.
      This allows to use the TFTModel without having to pass future_covariates to `fit()` and `train()`.
      It gives a value to the position of each step from input and output chunk relative to the prediction
      point. The values are normalized with `input_chunk_length`.
  likelihood
      The likelihood model to be used for probabilistic forecasts. By default, the TFT uses
      a ``QuantileRegression`` likelihood.
  norm_type: str | nn.Module
      The type of LayerNorm variant to use.
  **kwargs
      all parameters required for :class:`darts.model.forecasting_models.PLForecastingModule` base class.
docstr_tok: ["PyTorch","module","implementing","the","TFT","architecture","from","`","this","paper","<","https",":","\/\/arxiv.org\/pdf\/1912.09363.pdf",">","`","_","The","implementation","is","built","upon","`","pytorch-forecasting","'s","TemporalFusionTransformer","<","https",":","\/\/pytorch-forecasting.readthedocs.io\/en\/latest\/models.html",">","`","_",".","Parameters","--","--","--","--","--","output_dim",":","Tuple","[","int",",","int","]","shape","of","output","given","by","(","n_targets",",","loss_size",")",".","(","loss_size","corresponds","to","nr_params","in","other","models",")",".","variables_meta",":","Dict","[","str",",","Dict","[","str",",","List","[","str","]","]","]","dict","containing","variable","encoder",",","decoder","variable","names","for","mapping","tensors","in","`","_TFTModule.forward","(",")","`","num_static_components","the","number","of","static","components","(","not","variables",")","of","the","input","target","series",".","This","is","either","equal","to","the","number","of","target","components","or","1.","hidden_size",":","int","hidden","state","size","of","the","TFT",".","It","is","the","main","hyper-parameter","and","common","across","the","internal","TFT","architecture",".","lstm_layers",":","int","number","of","layers","for","the","Long","Short","Term","Memory","(","LSTM",")","Encoder","and","Decoder","(","1","is","a","good","default",")",".","num_attention_heads",":","int","number","of","attention","heads","(","4","is","a","good","default",")","full_attention",":","bool","If","`","True","`",",","applies","multi-head","attention","query","on","past","(","encoder",")","and","future","(","decoder",")","parts",".","Otherwise",",","only","queries","on","future","part",".","Defaults","to","`","False","`",".","feed_forward","Set","the","feedforward","network","block",".","default","`","GatedResidualNetwork","`","or","one","of","the","glu","variant",".","Defaults","to","`","GatedResidualNetwork","`",".","hidden_continuous_size",":","int","default","for","hidden","size","for","processing","continuous","variables",".","categorical_embedding_sizes",":","dict","A","dictionary","containing","embedding","sizes","for","categorical","static","covariates",".","The","keys","are","the","column","names","of","the","categorical","static","covariates",".","The","values","are","tuples","of","integers","with","`","(","number","of","unique","categories",",","embedding","size",")","`",".","For","example","`","{","``","some_column","''",":","(","64",",","8",")","}","`",".","Note","that","`","TorchForecastingModels","`","can","only","handle","numeric","data",".","Consider","transforming\/encoding","your","data","with","`","darts.dataprocessing.transformers.static_covariates_transformer.StaticCovariatesTransformer","`",".","dropout",":","float","Fraction","of","neurons","affected","by","Dropout",".","add_relative_index",":","bool","Whether","to","add","positional","values","to","future","covariates",".","Defaults","to","`","False","`",".","This","allows","to","use","the","TFTModel","without","having","to","pass","future_covariates","to","`","fit","(",")","`","and","`","train","(",")","`",".","It","gives","a","value","to","the","position","of","each","step","from","input","and","output","chunk","relative","to","the","prediction","point",".","The","values","are","normalized","with","`","input_chunk_length","`",".","likelihood","The","likelihood","model","to","be","used","for","probabilistic","forecasts",".","By","default",",","the","TFT","uses","a","``","QuantileRegression","``","likelihood",".","norm_type",":","str","|","nn.Module","The","type","of","LayerNorm","variant","to","use",".","*","*","kwargs","all","parameters","required","for",":","class",":","`","darts.model.forecasting_models.PLForecastingModule","`","base","class","."]
code: def __init__(
      self,
      output_dim: Tuple[int, int],
      variables_meta: Dict[str, Dict[str, List[str]]],
      num_static_components: int,
      hidden_size: Union[int, List[int]],
      lstm_layers: int,
      num_attention_heads: int,
      full_attention: bool,
      feed_forward: str,
      hidden_continuous_size: int,
      categorical_embedding_sizes: Dict[str, Tuple[int, int]],
      dropout: float,
      add_relative_index: bool,
      norm_type: Union[str, nn.Module],
      **kwargs,
  ):
      """PyTorch module implementing the TFT architecture from `this paper <https://arxiv.org/pdf/1912.09363.pdf>`_
      The implementation is built upon `pytorch-forecasting's TemporalFusionTransformer
      <https://pytorch-forecasting.readthedocs.io/en/latest/models.html>`_.
  
      Parameters
      ----------
      output_dim : Tuple[int, int]
          shape of output given by (n_targets, loss_size). (loss_size corresponds to nr_params in other models).
      variables_meta : Dict[str, Dict[str, List[str]]]
          dict containing variable encoder, decoder variable names for mapping tensors in `_TFTModule.forward()`
      num_static_components
          the number of static components (not variables) of the input target series. This is either equal to the
          number of target components or 1.
      hidden_size : int
          hidden state size of the TFT. It is the main hyper-parameter and common across the internal TFT
          architecture.
      lstm_layers : int
          number of layers for the Long Short Term Memory (LSTM) Encoder and Decoder (1 is a good default).
      num_attention_heads : int
          number of attention heads (4 is a good default)
      full_attention : bool
          If `True`, applies multi-head attention query on past (encoder) and future (decoder) parts. Otherwise,
          only queries on future part. Defaults to `False`.
      feed_forward
          Set the feedforward network block. default `GatedResidualNetwork` or one of the  glu variant.
          Defaults to `GatedResidualNetwork`.
      hidden_continuous_size : int
          default for hidden size for processing continuous variables.
      categorical_embedding_sizes : dict
          A dictionary containing embedding sizes for categorical static covariates. The keys are the column names
          of the categorical static covariates. The values are tuples of integers with
          `(number of unique categories, embedding size)`. For example `{"some_column": (64, 8)}`.
          Note that `TorchForecastingModels` can only handle numeric data. Consider transforming/encoding your data
          with `darts.dataprocessing.transformers.static_covariates_transformer.StaticCovariatesTransformer`.
      dropout : float
          Fraction of neurons affected by Dropout.
      add_relative_index : bool
          Whether to add positional values to future covariates. Defaults to `False`.
          This allows to use the TFTModel without having to pass future_covariates to `fit()` and `train()`.
          It gives a value to the position of each step from input and output chunk relative to the prediction
          point. The values are normalized with `input_chunk_length`.
      likelihood
          The likelihood model to be used for probabilistic forecasts. By default, the TFT uses
          a ``QuantileRegression`` likelihood.
      norm_type: str | nn.Module
          The type of LayerNorm variant to use.
      **kwargs
          all parameters required for :class:`darts.model.forecasting_models.PLForecastingModule` base class.
      """
  
      super().__init__(**kwargs)
  
      self.n_targets, self.loss_size = output_dim
      self.variables_meta = variables_meta
      self.num_static_components = num_static_components
      self.hidden_size = hidden_size
      self.hidden_continuous_size = hidden_continuous_size
      self.categorical_embedding_sizes = categorical_embedding_sizes
      self.lstm_layers = lstm_layers
      self.num_attention_heads = num_attention_heads
      self.full_attention = full_attention
      self.feed_forward = feed_forward
      self.dropout = dropout
      self.add_relative_index = add_relative_index
  
      if isinstance(norm_type, str):
          try:
              self.layer_norm = getattr(layer_norm_variants, norm_type)
          except AttributeError:
              raise_log(
                  AttributeError(
                      "please provide a valid layer norm type"
                  ),
              )
      else:
          self.layer_norm = norm_type
  
      # initialize last batch size to check if new mask needs to be generated
      self.batch_size_last = -1
      self.attention_mask = None
      self.relative_index = None
  
      # general information on variable name endings:
      # _vsn: VariableSelectionNetwork
      # _grn: GatedResidualNetwork
      # _glu: GatedLinearUnit
      # _gan: GateAddNorm
      # _attn: Attention
  
      # # processing inputs
      # embeddings
      self.input_embeddings = _MultiEmbedding(
          embedding_sizes=categorical_embedding_sizes,
          variable_names=self.categorical_static_variables,
      )
  
      # continuous variable processing
      self.prescalers_linear = {
          name: nn.Linear(
              1
              if name not in self.numeric_static_variables
              else self.num_static_components,
              self.hidden_continuous_size,
          )
          for name in self.reals
      }
  
      # static (categorical and numerical) variables
      static_input_sizes = {
          name: self.input_embeddings.output_size[name]
          for name in self.categorical_static_variables
      }
      static_input_sizes.update(
          {
              name: self.hidden_continuous_size
              for name in self.numeric_static_variables
          }
      )
  
      self.static_covariates_vsn = _VariableSelectionNetwork(
          input_sizes=static_input_sizes,
          hidden_size=self.hidden_size,
          input_embedding_flags={
              name: True for name in self.categorical_static_variables
          },
          dropout=self.dropout,
          prescalers=self.prescalers_linear,
          single_variable_grns={},
          context_size=None,  # no context for static variables
          layer_norm=self.layer_norm,
      )
  
      # variable selection for encoder and decoder
      encoder_input_sizes = {
          name: self.hidden_continuous_size
          for name in self.encoder_variables
      }
  
      decoder_input_sizes = {
          name: self.hidden_continuous_size
          for name in self.decoder_variables
      }
  
      self.encoder_vsn = _VariableSelectionNetwork(
          input_sizes=encoder_input_sizes,
          hidden_size=self.hidden_size,
          input_embedding_flags={},  # this would be required for non-static categorical inputs
          dropout=self.dropout,
          context_size=self.hidden_size,
          prescalers=self.prescalers_linear,
          single_variable_grns={},
          layer_norm=self.layer_norm,
      )
  
      self.decoder_vsn = _VariableSelectionNetwork(
          input_sizes=decoder_input_sizes,
          hidden_size=self.hidden_size,
          input_embedding_flags={},  # this would be required for non-static categorical inputs
          dropout=self.dropout,
          context_size=self.hidden_size,
          prescalers=self.prescalers_linear,
          single_variable_grns={},
          layer_norm=self.layer_norm,
      )
  
      # static encoders
      # for variable selection
      self.static_context_grn = _GatedResidualNetwork(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          output_size=self.hidden_size,
          dropout=self.dropout,
          layer_norm=self.layer_norm,
      )
  
      # for hidden state of the lstm
      self.static_context_hidden_encoder_grn = _GatedResidualNetwork(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          output_size=self.hidden_size,
          dropout=self.dropout,
          layer_norm=self.layer_norm,
      )
  
      # for cell state of the lstm
      self.static_context_cell_encoder_grn = _GatedResidualNetwork(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          output_size=self.hidden_size,
          dropout=self.dropout,
          layer_norm=self.layer_norm,
      )
  
      # for post lstm static enrichment
      self.static_context_enrichment = _GatedResidualNetwork(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          output_size=self.hidden_size,
          dropout=self.dropout,
          layer_norm=self.layer_norm,
      )
  
      # lstm encoder (history) and decoder (future) for local processing
      self.lstm_encoder = _LSTM(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          num_layers=self.lstm_layers,
          dropout=self.dropout if self.lstm_layers > 1 else 0,
          batch_first=True,
      )
  
      self.lstm_decoder = _LSTM(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          num_layers=self.lstm_layers,
          dropout=self.dropout if self.lstm_layers > 1 else 0,
          batch_first=True,
      )
  
      # post lstm GateAddNorm
      self.post_lstm_gan = _GateAddNorm(
          input_size=self.hidden_size,
          dropout=dropout,
          layer_norm=self.layer_norm,
      )
  
      # static enrichment and processing past LSTM
      self.static_enrichment_grn = _GatedResidualNetwork(
          input_size=self.hidden_size,
          hidden_size=self.hidden_size,
          output_size=self.hidden_size,
          dropout=self.dropout,
          context_size=self.hidden_size,
          layer_norm=self.layer_norm,
      )
  
      # attention for long-range processing
      self.multihead_attn = _InterpretableMultiHeadAttention(
          d_model=self.hidden_size,
          n_head=self.num_attention_heads,
          dropout=self.dropout,
      )
      self.post_attn_gan = _GateAddNorm(
          self.hidden_size,
          dropout=self.dropout,
          layer_norm=self.layer_norm,
      )
  
      if self.feed_forward == "GatedResidualNetwork":
          self.feed_forward_block = _GatedResidualNetwork(
              self.hidden_size,
              self.hidden_size,
              self.hidden_size,
              dropout=self.dropout,
              layer_norm=self.layer_norm,
          )
      else:
          raise_if_not(
              self.feed_forward in GLU_FFN,
              f"'{self.feed_forward}' is not in {GLU_FFN + ['GatedResidualNetwork']}",
          )
          # use glu variant feedforward layers
          # 4 is a commonly used feedforward multiplier
          self.feed_forward_block = getattr(
              glu_variants, self.feed_forward
          )(
              d_model=self.hidden_size,
              d_ff=self.hidden_size * 4,
              dropout=dropout,
          )
  
      # output processing -> no dropout at this late stage
      self.pre_output_gan = _GateAddNorm(
          self.hidden_size, dropout=None, layer_norm=self.layer_norm
      )
  
      self.output_layer = nn.Linear(
          self.hidden_size, self.n_targets * self.loss_size
      )
  
      self._attn_out_weights = None
      self._static_covariate_var = None
      self._encoder_sparse_weights = None
      self._decoder_sparse_weights = None
  
imports: from typing import Dict, List, Optional, Sequence, Tuple, Union
  import numpy
  import pandas
  import torch
  from torch import nn
  from torch.nn import LSTM
  from darts import TimeSeries
  from darts.logging import get_logger, raise_if, raise_if_not, raise_log
  from darts.models.components import glu_variants, layer_norm_variants
  from darts.models.components.glu_variants import GLU_FFN
  from darts.models.forecasting.pl_forecasting_module import PLMixedCovariatesModule, io_processor
  from darts.models.forecasting.tft_submodels import _GateAddNorm, _GatedResidualNetwork, _InterpretableMultiHeadAttention, _MultiEmbedding, _VariableSelectionNetwork, get_embedding_size
  from darts.models.forecasting.torch_forecasting_model import MixedCovariatesTorchModel
  from darts.utils.data import MixedCovariatesSequentialDataset, MixedCovariatesTrainingDataset, TrainingDataset
  from darts.utils.likelihood_models import Likelihood, QuantileRegression
label_desc: Use image node_id 1 to create a new _TFTModule object from inherited base classes: PLMixedCovariatesModule with example: obj = _TFTModule(output_dim, variables_meta, num_static_components, hidden_size, lstm_layers, num_attention_heads, full_attention, feed_forward, hidden_continuous_size, categorical_embedding_sizes, dropout, add_relative_index, norm_type)