docstr: Temporal Fusion Transformer for forecasting timeseries - use its :py:meth:`~from_dataset` method if possible.
  
  Implementation of the article
  `Temporal Fusion Transformers for Interpretable Multi-horizon Time Series
  Forecasting <https://arxiv.org/pdf/1912.09363.pdf>`_. The network outperforms DeepAR by Amazon by 36-69%
  in benchmarks.
  
  Enhancements compared to the original implementation (apart from capabilities added through base model
  such as monotone constraints):
  
  * static variables can be continuous
  * multiple categorical variables can be summarized with an EmbeddingBag
  * variable encoder and decoder length by sample
  * categorical embeddings are not transformed by variable selection network (because it is a redundant operation)
  * variable dimension in variable selection network are scaled up via linear interpolation to reduce
    number of parameters
  * non-linear variable processing in variable selection network can be shared among decoder and encoder
    (not shared by default)
  
  Tune its hyperparameters with
  :py:func:`~pytorch_forecasting.models.temporal_fusion_transformer.tuning.optimize_hyperparameters`.
  
  Args:
  
      hidden_size: hidden size of network which is its main hyperparameter and can range from 8 to 512
      lstm_layers: number of LSTM layers (2 is mostly optimal)
      dropout: dropout rate
      output_size: number of outputs (e.g. number of quantiles for QuantileLoss and one target or list
          of output sizes).
      loss: loss function taking prediction and targets
      attention_head_size: number of attention heads (4 is a good default)
      max_encoder_length: length to encode (can be far longer than the decoder length but does not have to be)
      static_categoricals: names of static categorical variables
      static_reals: names of static continuous variables
      time_varying_categoricals_encoder: names of categorical variables for encoder
      time_varying_categoricals_decoder: names of categorical variables for decoder
      time_varying_reals_encoder: names of continuous variables for encoder
      time_varying_reals_decoder: names of continuous variables for decoder
      categorical_groups: dictionary where values
          are list of categorical variables that are forming together a new categorical
          variable which is the key in the dictionary
      x_reals: order of continuous variables in tensor passed to forward function
      x_categoricals: order of categorical variables in tensor passed to forward function
      hidden_continuous_size: default for hidden size for processing continous variables (similar to categorical
          embedding size)
      hidden_continuous_sizes: dictionary mapping continuous input indices to sizes for variable selection
          (fallback to hidden_continuous_size if index is not in dictionary)
      embedding_sizes: dictionary mapping (string) indices to tuple of number of categorical classes and
          embedding size
      embedding_paddings: list of indices for embeddings which transform the zero's embedding to a zero vector
      embedding_labels: dictionary mapping (string) indices to list of categorical labels
      learning_rate: learning rate
      log_interval: log predictions every x batches, do not log if 0 or less, log interpretation if > 0. If < 1.0
          , will log multiple entries per batch. Defaults to -1.
      log_val_interval: frequency with which to log validation set metrics, defaults to log_interval
      log_gradient_flow: if to log gradient flow, this takes time and should be only done to diagnose training
          failures
      reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10
      monotone_constaints (Dict[str, int]): dictionary of monotonicity constraints for continuous decoder
          variables mapping
          position (e.g. ``"0"`` for first position) to constraint (``-1`` for negative and ``+1`` for positive,
          larger numbers add more weight to the constraint vs. the loss but are usually not necessary).
          This constraint significantly slows down training. Defaults to {}.
      share_single_variable_networks (bool): if to share the single variable networks between the encoder and
          decoder. Defaults to False.
      causal_attention (bool): If to attend only at previous timesteps in the decoder or also include future
          predictions. Defaults to True.
      logging_metrics (nn.ModuleList[LightningMetric]): list of metrics that are logged during training.
          Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]).
      **kwargs: additional arguments to :py:class:`~BaseModel`.
docstr_tok: ["Temporal","Fusion","Transformer","for","forecasting","timeseries","-","use","its",":","py",":","meth",":","`","~from_dataset","`","method","if","possible",".","Implementation","of","the","article","`","Temporal","Fusion","Transformers","for","Interpretable","Multi-horizon","Time","Series","Forecasting","<","https",":","\/\/arxiv.org\/pdf\/1912.09363.pdf",">","`","_",".","The","network","outperforms","DeepAR","by","Amazon","by","36-69","%","in","benchmarks",".","Enhancements","compared","to","the","original","implementation","(","apart","from","capabilities","added","through","base","model","such","as","monotone","constraints",")",":","*","static","variables","can","be","continuous","*","multiple","categorical","variables","can","be","summarized","with","an","EmbeddingBag","*","variable","encoder","and","decoder","length","by","sample","*","categorical","embeddings","are","not","transformed","by","variable","selection","network","(","because","it","is","a","redundant","operation",")","*","variable","dimension","in","variable","selection","network","are","scaled","up","via","linear","interpolation","to","reduce","number","of","parameters","*","non-linear","variable","processing","in","variable","selection","network","can","be","shared","among","decoder","and","encoder","(","not","shared","by","default",")","Tune","its","hyperparameters","with",":","py",":","func",":","`","~pytorch_forecasting.models.temporal_fusion_transformer.tuning.optimize_hyperparameters","`",".","Args",":","hidden_size",":","hidden","size","of","network","which","is","its","main","hyperparameter","and","can","range","from","8","to","512","lstm_layers",":","number","of","LSTM","layers","(","2","is","mostly","optimal",")","dropout",":","dropout","rate","output_size",":","number","of","outputs","(","e.g",".","number","of","quantiles","for","QuantileLoss","and","one","target","or","list","of","output","sizes",")",".","loss",":","loss","function","taking","prediction","and","targets","attention_head_size",":","number","of","attention","heads","(","4","is","a","good","default",")","max_encoder_length",":","length","to","encode","(","can","be","far","longer","than","the","decoder","length","but","does","not","have","to","be",")","static_categoricals",":","names","of","static","categorical","variables","static_reals",":","names","of","static","continuous","variables","time_varying_categoricals_encoder",":","names","of","categorical","variables","for","encoder","time_varying_categoricals_decoder",":","names","of","categorical","variables","for","decoder","time_varying_reals_encoder",":","names","of","continuous","variables","for","encoder","time_varying_reals_decoder",":","names","of","continuous","variables","for","decoder","categorical_groups",":","dictionary","where","values","are","list","of","categorical","variables","that","are","forming","together","a","new","categorical","variable","which","is","the","key","in","the","dictionary","x_reals",":","order","of","continuous","variables","in","tensor","passed","to","forward","function","x_categoricals",":","order","of","categorical","variables","in","tensor","passed","to","forward","function","hidden_continuous_size",":","default","for","hidden","size","for","processing","continous","variables","(","similar","to","categorical","embedding","size",")","hidden_continuous_sizes",":","dictionary","mapping","continuous","input","indices","to","sizes","for","variable","selection","(","fallback","to","hidden_continuous_size","if","index","is","not","in","dictionary",")","embedding_sizes",":","dictionary","mapping","(","string",")","indices","to","tuple","of","number","of","categorical","classes","and","embedding","size","embedding_paddings",":","list","of","indices","for","embeddings","which","transform","the","zero","'s","embedding","to","a","zero","vector","embedding_labels",":","dictionary","mapping","(","string",")","indices","to","list","of","categorical","labels","learning_rate",":","learning","rate","log_interval",":","log","predictions","every","x","batches",",","do","not","log","if","0","or","less",",","log","interpretation","if",">","0",".","If","<","1.0",",","will","log","multiple","entries","per","batch",".","Defaults","to","-1.","log_val_interval",":","frequency","with","which","to","log","validation","set","metrics",",","defaults","to","log_interval","log_gradient_flow",":","if","to","log","gradient","flow",",","this","takes","time","and","should","be","only","done","to","diagnose","training","failures","reduce_on_plateau_patience","(","int",")",":","patience","after","which","learning","rate","is","reduced","by","a","factor","of","10","monotone_constaints","(","Dict","[","str",",","int","]",")",":","dictionary","of","monotonicity","constraints","for","continuous","decoder","variables","mapping","position","(","e.g",".","``","``","0","''","``","for","first","position",")","to","constraint","(","``","-1","``","for","negative","and","``","+1","``","for","positive",",","larger","numbers","add","more","weight","to","the","constraint","vs.","the","loss","but","are","usually","not","necessary",")",".","This","constraint","significantly","slows","down","training",".","Defaults","to","{","}",".","share_single_variable_networks","(","bool",")",":","if","to","share","the","single","variable","networks","between","the","encoder","and","decoder",".","Defaults","to","False",".","causal_attention","(","bool",")",":","If","to","attend","only","at","previous","timesteps","in","the","decoder","or","also","include","future","predictions",".","Defaults","to","True",".","logging_metrics","(","nn.ModuleList","[","LightningMetric","]",")",":","list","of","metrics","that","are","logged","during","training",".","Defaults","to","nn.ModuleList","(","[","SMAPE","(",")",",","MAE","(",")",",","RMSE","(",")",",","MAPE","(",")","]",")",".","*","*","kwargs",":","additional","arguments","to",":","py",":","class",":","`","~BaseModel","`","."]
code: def __init__(
      self,
      hidden_size: int = 16,
      lstm_layers: int = 1,
      dropout: float = 0.1,
      output_size: Union[int, List[int]] = 7,
      loss: MultiHorizonMetric = None,
      attention_head_size: int = 4,
      max_encoder_length: int = 10,
      static_categoricals: List[str] = [],
      static_reals: List[str] = [],
      time_varying_categoricals_encoder: List[str] = [],
      time_varying_categoricals_decoder: List[str] = [],
      categorical_groups: Dict[str, List[str]] = {},
      time_varying_reals_encoder: List[str] = [],
      time_varying_reals_decoder: List[str] = [],
      x_reals: List[str] = [],
      x_categoricals: List[str] = [],
      hidden_continuous_size: int = 8,
      hidden_continuous_sizes: Dict[str, int] = {},
      embedding_sizes: Dict[str, Tuple[int, int]] = {},
      embedding_paddings: List[str] = [],
      embedding_labels: Dict[str, np.ndarray] = {},
      learning_rate: float = 1e-3,
      log_interval: Union[int, float] = -1,
      log_val_interval: Union[int, float] = None,
      log_gradient_flow: bool = False,
      reduce_on_plateau_patience: int = 1000,
      monotone_constaints: Dict[str, int] = {},
      share_single_variable_networks: bool = False,
      causal_attention: bool = True,
      logging_metrics: nn.ModuleList = None,
      **kwargs,
  ):
      """
      Temporal Fusion Transformer for forecasting timeseries - use its :py:meth:`~from_dataset` method if possible.
  
      Implementation of the article
      `Temporal Fusion Transformers for Interpretable Multi-horizon Time Series
      Forecasting <https://arxiv.org/pdf/1912.09363.pdf>`_. The network outperforms DeepAR by Amazon by 36-69%
      in benchmarks.
  
      Enhancements compared to the original implementation (apart from capabilities added through base model
      such as monotone constraints):
  
      * static variables can be continuous
      * multiple categorical variables can be summarized with an EmbeddingBag
      * variable encoder and decoder length by sample
      * categorical embeddings are not transformed by variable selection network (because it is a redundant operation)
      * variable dimension in variable selection network are scaled up via linear interpolation to reduce
        number of parameters
      * non-linear variable processing in variable selection network can be shared among decoder and encoder
        (not shared by default)
  
      Tune its hyperparameters with
      :py:func:`~pytorch_forecasting.models.temporal_fusion_transformer.tuning.optimize_hyperparameters`.
  
      Args:
  
          hidden_size: hidden size of network which is its main hyperparameter and can range from 8 to 512
          lstm_layers: number of LSTM layers (2 is mostly optimal)
          dropout: dropout rate
          output_size: number of outputs (e.g. number of quantiles for QuantileLoss and one target or list
              of output sizes).
          loss: loss function taking prediction and targets
          attention_head_size: number of attention heads (4 is a good default)
          max_encoder_length: length to encode (can be far longer than the decoder length but does not have to be)
          static_categoricals: names of static categorical variables
          static_reals: names of static continuous variables
          time_varying_categoricals_encoder: names of categorical variables for encoder
          time_varying_categoricals_decoder: names of categorical variables for decoder
          time_varying_reals_encoder: names of continuous variables for encoder
          time_varying_reals_decoder: names of continuous variables for decoder
          categorical_groups: dictionary where values
              are list of categorical variables that are forming together a new categorical
              variable which is the key in the dictionary
          x_reals: order of continuous variables in tensor passed to forward function
          x_categoricals: order of categorical variables in tensor passed to forward function
          hidden_continuous_size: default for hidden size for processing continous variables (similar to categorical
              embedding size)
          hidden_continuous_sizes: dictionary mapping continuous input indices to sizes for variable selection
              (fallback to hidden_continuous_size if index is not in dictionary)
          embedding_sizes: dictionary mapping (string) indices to tuple of number of categorical classes and
              embedding size
          embedding_paddings: list of indices for embeddings which transform the zero's embedding to a zero vector
          embedding_labels: dictionary mapping (string) indices to list of categorical labels
          learning_rate: learning rate
          log_interval: log predictions every x batches, do not log if 0 or less, log interpretation if > 0. If < 1.0
              , will log multiple entries per batch. Defaults to -1.
          log_val_interval: frequency with which to log validation set metrics, defaults to log_interval
          log_gradient_flow: if to log gradient flow, this takes time and should be only done to diagnose training
              failures
          reduce_on_plateau_patience (int): patience after which learning rate is reduced by a factor of 10
          monotone_constaints (Dict[str, int]): dictionary of monotonicity constraints for continuous decoder
              variables mapping
              position (e.g. ``"0"`` for first position) to constraint (``-1`` for negative and ``+1`` for positive,
              larger numbers add more weight to the constraint vs. the loss but are usually not necessary).
              This constraint significantly slows down training. Defaults to {}.
          share_single_variable_networks (bool): if to share the single variable networks between the encoder and
              decoder. Defaults to False.
          causal_attention (bool): If to attend only at previous timesteps in the decoder or also include future
              predictions. Defaults to True.
          logging_metrics (nn.ModuleList[LightningMetric]): list of metrics that are logged during training.
              Defaults to nn.ModuleList([SMAPE(), MAE(), RMSE(), MAPE()]).
          **kwargs: additional arguments to :py:class:`~BaseModel`.
      """
      if logging_metrics is None:
          logging_metrics = nn.ModuleList(
              [SMAPE(), MAE(), RMSE(), MAPE()]
          )
      if loss is None:
          loss = QuantileLoss()
      self.save_hyperparameters()
      # store loss function separately as it is a module
      assert isinstance(
          loss, LightningMetric
      ), "Loss has to be a PyTorch Lightning `Metric`"
      super().__init__(
          loss=loss, logging_metrics=logging_metrics, **kwargs
      )
  
      # processing inputs
      # embeddings
      self.input_embeddings = MultiEmbedding(
          embedding_sizes=self.hparams.embedding_sizes,
          categorical_groups=self.hparams.categorical_groups,
          embedding_paddings=self.hparams.embedding_paddings,
          x_categoricals=self.hparams.x_categoricals,
          max_embedding_size=self.hparams.hidden_size,
      )
  
      # continuous variable processing
      self.prescalers = nn.ModuleDict(
          {
              name: nn.Linear(
                  1,
                  self.hparams.hidden_continuous_sizes.get(
                      name, self.hparams.hidden_continuous_size
                  ),
              )
              for name in self.reals
          }
      )
  
      # variable selection
      # variable selection for static variables
      static_input_sizes = {
          name: self.input_embeddings.output_size[name]
          for name in self.hparams.static_categoricals
      }
      static_input_sizes.update(
          {
              name: self.hparams.hidden_continuous_sizes.get(
                  name, self.hparams.hidden_continuous_size
              )
              for name in self.hparams.static_reals
          }
      )
      self.static_variable_selection = VariableSelectionNetwork(
          input_sizes=static_input_sizes,
          hidden_size=self.hparams.hidden_size,
          input_embedding_flags={
              name: True for name in self.hparams.static_categoricals
          },
          dropout=self.hparams.dropout,
          prescalers=self.prescalers,
      )
  
      # variable selection for encoder and decoder
      encoder_input_sizes = {
          name: self.input_embeddings.output_size[name]
          for name in self.hparams.time_varying_categoricals_encoder
      }
      encoder_input_sizes.update(
          {
              name: self.hparams.hidden_continuous_sizes.get(
                  name, self.hparams.hidden_continuous_size
              )
              for name in self.hparams.time_varying_reals_encoder
          }
      )
  
      decoder_input_sizes = {
          name: self.input_embeddings.output_size[name]
          for name in self.hparams.time_varying_categoricals_decoder
      }
      decoder_input_sizes.update(
          {
              name: self.hparams.hidden_continuous_sizes.get(
                  name, self.hparams.hidden_continuous_size
              )
              for name in self.hparams.time_varying_reals_decoder
          }
      )
  
      # create single variable grns that are shared across decoder and encoder
      if self.hparams.share_single_variable_networks:
          self.shared_single_variable_grns = nn.ModuleDict()
          for name, input_size in encoder_input_sizes.items():
              self.shared_single_variable_grns[
                  name
              ] = GatedResidualNetwork(
                  input_size,
                  min(input_size, self.hparams.hidden_size),
                  self.hparams.hidden_size,
                  self.hparams.dropout,
              )
          for name, input_size in decoder_input_sizes.items():
              if name not in self.shared_single_variable_grns:
                  self.shared_single_variable_grns[
                      name
                  ] = GatedResidualNetwork(
                      input_size,
                      min(input_size, self.hparams.hidden_size),
                      self.hparams.hidden_size,
                      self.hparams.dropout,
                  )
  
      self.encoder_variable_selection = VariableSelectionNetwork(
          input_sizes=encoder_input_sizes,
          hidden_size=self.hparams.hidden_size,
          input_embedding_flags={
              name: True
              for name in self.hparams.time_varying_categoricals_encoder
          },
          dropout=self.hparams.dropout,
          context_size=self.hparams.hidden_size,
          prescalers=self.prescalers,
          single_variable_grns={}
          if not self.hparams.share_single_variable_networks
          else self.shared_single_variable_grns,
      )
  
      self.decoder_variable_selection = VariableSelectionNetwork(
          input_sizes=decoder_input_sizes,
          hidden_size=self.hparams.hidden_size,
          input_embedding_flags={
              name: True
              for name in self.hparams.time_varying_categoricals_decoder
          },
          dropout=self.hparams.dropout,
          context_size=self.hparams.hidden_size,
          prescalers=self.prescalers,
          single_variable_grns={}
          if not self.hparams.share_single_variable_networks
          else self.shared_single_variable_grns,
      )
  
      # static encoders
      # for variable selection
      self.static_context_variable_selection = GatedResidualNetwork(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          output_size=self.hparams.hidden_size,
          dropout=self.hparams.dropout,
      )
  
      # for hidden state of the lstm
      self.static_context_initial_hidden_lstm = GatedResidualNetwork(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          output_size=self.hparams.hidden_size,
          dropout=self.hparams.dropout,
      )
  
      # for cell state of the lstm
      self.static_context_initial_cell_lstm = GatedResidualNetwork(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          output_size=self.hparams.hidden_size,
          dropout=self.hparams.dropout,
      )
  
      # for post lstm static enrichment
      self.static_context_enrichment = GatedResidualNetwork(
          self.hparams.hidden_size,
          self.hparams.hidden_size,
          self.hparams.hidden_size,
          self.hparams.dropout,
      )
  
      # lstm encoder (history) and decoder (future) for local processing
      self.lstm_encoder = LSTM(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          num_layers=self.hparams.lstm_layers,
          dropout=self.hparams.dropout
          if self.hparams.lstm_layers > 1
          else 0,
          batch_first=True,
      )
  
      self.lstm_decoder = LSTM(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          num_layers=self.hparams.lstm_layers,
          dropout=self.hparams.dropout
          if self.hparams.lstm_layers > 1
          else 0,
          batch_first=True,
      )
  
      # skip connection for lstm
      self.post_lstm_gate_encoder = GatedLinearUnit(
          self.hparams.hidden_size, dropout=self.hparams.dropout
      )
      self.post_lstm_gate_decoder = self.post_lstm_gate_encoder
      # self.post_lstm_gate_decoder = GatedLinearUnit(self.hparams.hidden_size, dropout=self.hparams.dropout)
      self.post_lstm_add_norm_encoder = AddNorm(
          self.hparams.hidden_size, trainable_add=False
      )
      # self.post_lstm_add_norm_decoder = AddNorm(self.hparams.hidden_size, trainable_add=True)
      self.post_lstm_add_norm_decoder = self.post_lstm_add_norm_encoder
  
      # static enrichment and processing past LSTM
      self.static_enrichment = GatedResidualNetwork(
          input_size=self.hparams.hidden_size,
          hidden_size=self.hparams.hidden_size,
          output_size=self.hparams.hidden_size,
          dropout=self.hparams.dropout,
          context_size=self.hparams.hidden_size,
      )
  
      # attention for long-range processing
      self.multihead_attn = InterpretableMultiHeadAttention(
          d_model=self.hparams.hidden_size,
          n_head=self.hparams.attention_head_size,
          dropout=self.hparams.dropout,
      )
      self.post_attn_gate_norm = GateAddNorm(
          self.hparams.hidden_size,
          dropout=self.hparams.dropout,
          trainable_add=False,
      )
      self.pos_wise_ff = GatedResidualNetwork(
          self.hparams.hidden_size,
          self.hparams.hidden_size,
          self.hparams.hidden_size,
          dropout=self.hparams.dropout,
      )
  
      # output processing -> no dropout at this late stage
      self.pre_output_gate_norm = GateAddNorm(
          self.hparams.hidden_size, dropout=None, trainable_add=False
      )
  
      if self.n_targets > 1:  # if to run with multiple targets
          self.output_layer = nn.ModuleList(
              [
                  nn.Linear(self.hparams.hidden_size, output_size)
                  for output_size in self.hparams.output_size
              ]
          )
      else:
          self.output_layer = nn.Linear(
              self.hparams.hidden_size, self.hparams.output_size
          )
  
imports: from copy import copy
  from typing import Dict, List, Tuple, Union
  from matplotlib import pyplot
  import numpy
  import torch
  from torch import nn
  from torchmetrics import Metric
  from pytorch_forecasting.data import TimeSeriesDataSet
  from pytorch_forecasting.data.encoders import NaNLabelEncoder
  from pytorch_forecasting.metrics import MAE, MAPE, MASE, RMSE, SMAPE, MultiHorizonMetric, MultiLoss, QuantileLoss
  from pytorch_forecasting.models.base_model import BaseModelWithCovariates
  from pytorch_forecasting.models.nn import LSTM, MultiEmbedding
  from pytorch_forecasting.models.temporal_fusion_transformer.sub_modules import AddNorm, GateAddNorm, GatedLinearUnit, GatedResidualNetwork, InterpretableMultiHeadAttention, VariableSelectionNetwork
  from pytorch_forecasting.utils import create_mask, detach, integer_histogram, masked_op, padded_stack, to_list
label_desc: Use image node_id 1 to create a new TemporalFusionTransformer object from inherited base classes: BaseModelWithCovariates with example: obj = TemporalFusionTransformer(hidden_size, lstm_layers, dropout, output_size, loss, attention_head_size, max_encoder_length, static_categoricals, static_reals, time_varying_categoricals_encoder, time_varying_categoricals_decoder, categorical_groups, time_varying_reals_encoder, time_varying_reals_decoder, x_reals, x_categoricals, hidden_continuous_size, hidden_continuous_sizes, embedding_sizes, embedding_paddings, embedding_labels, learning_rate, log_interval, log_val_interval, log_gradient_flow, reduce_on_plateau_patience, monotone_constaints, share_single_variable_networks, causal_attention, logging_metrics)