docstr: TFT model forward pass.
  
  Parameters
  ----------
  x_in
      comes as tuple `(x_past, x_future, x_static)` where `x_past` is the input/past chunk and `x_future`
      is the output/future chunk. Input dimensions are `(n_samples, n_time_steps, n_variables)`
  
  Returns
  -------
  torch.Tensor
      the output tensor
docstr_tok: ["TFT","model","forward","pass",".","Parameters","--","--","--","--","--","x_in","comes","as","tuple","`","(","x_past",",","x_future",",","x_static",")","`","where","`","x_past","`","is","the","input\/past","chunk","and","`","x_future","`","is","the","output\/future","chunk",".","Input","dimensions","are","`","(","n_samples",",","n_time_steps",",","n_variables",")","`","Returns","--","--","--","-","torch.Tensor","the","output","tensor"]
code: def forward(
      self,
      x_in: Tuple[
          torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]
      ],
  ) -> torch.Tensor:
      """TFT model forward pass.
  
      Parameters
      ----------
      x_in
          comes as tuple `(x_past, x_future, x_static)` where `x_past` is the input/past chunk and `x_future`
          is the output/future chunk. Input dimensions are `(n_samples, n_time_steps, n_variables)`
  
      Returns
      -------
      torch.Tensor
          the output tensor
      """
      x_cont_past, x_cont_future, x_static = x_in
      dim_samples, dim_time, dim_variable = 0, 1, 2
      device = x_in[0].device
  
      batch_size = x_cont_past.shape[dim_samples]
      encoder_length = self.input_chunk_length
      decoder_length = self.output_chunk_length
      time_steps = encoder_length + decoder_length
  
      # avoid unnecessary regeneration of attention mask
      if batch_size != self.batch_size_last:
          self.attention_mask = self.get_attention_mask_future(
              encoder_length=encoder_length,
              decoder_length=decoder_length,
              batch_size=batch_size,
              device=device,
              full_attention=self.full_attention,
          )
          if self.add_relative_index:
              self.relative_index = self.get_relative_index(
                  encoder_length=encoder_length,
                  decoder_length=decoder_length,
                  batch_size=batch_size,
                  device=device,
                  dtype=x_cont_past.dtype,
              )
  
          self.batch_size_last = batch_size
  
      if self.add_relative_index:
          x_cont_past = torch.cat(
              [
                  ts[:, :encoder_length, :]
                  for ts in [x_cont_past, self.relative_index]
                  if ts is not None
              ],
              dim=dim_variable,
          )
          x_cont_future = torch.cat(
              [
                  ts[:, -decoder_length:, :]
                  for ts in [x_cont_future, self.relative_index]
                  if ts is not None
              ],
              dim=dim_variable,
          )
  
      input_vectors_past = {
          name: x_cont_past[..., idx].unsqueeze(-1)
          for idx, name in enumerate(self.encoder_variables)
      }
      input_vectors_future = {
          name: x_cont_future[..., idx].unsqueeze(-1)
          for idx, name in enumerate(self.decoder_variables)
      }
  
      # Embedding and variable selection
      if self.static_variables:
          # categorical static covariate embeddings
          if self.categorical_static_variables:
              static_embedding = self.input_embeddings(
                  torch.cat(
                      [
                          x_static[:, :, idx]
                          for idx, name in enumerate(
                              self.static_variables
                          )
                          if name in self.categorical_static_variables
                      ],
                      dim=1,
                  ).int()
              )
          else:
              static_embedding = {}
          # add numerical static covariates
          static_embedding.update(
              {
                  name: x_static[:, :, idx]
                  for idx, name in enumerate(self.static_variables)
                  if name in self.numeric_static_variables
              }
          )
          (
              static_embedding,
              static_covariate_var,
          ) = self.static_covariates_vsn(static_embedding)
      else:
          static_embedding = torch.zeros(
              (x_cont_past.shape[0], self.hidden_size),
              dtype=x_cont_past.dtype,
              device=device,
          )
          static_covariate_var = None
  
      static_context_expanded = self.expand_static_context(
          context=self.static_context_grn(static_embedding),
          time_steps=time_steps,
      )
  
      embeddings_varying_encoder = {
          name: input_vectors_past[name]
          for name in self.encoder_variables
      }
      (
          embeddings_varying_encoder,
          encoder_sparse_weights,
      ) = self.encoder_vsn(
          x=embeddings_varying_encoder,
          context=static_context_expanded[:, :encoder_length],
      )
  
      embeddings_varying_decoder = {
          name: input_vectors_future[name]
          for name in self.decoder_variables
      }
      (
          embeddings_varying_decoder,
          decoder_sparse_weights,
      ) = self.decoder_vsn(
          x=embeddings_varying_decoder,
          context=static_context_expanded[:, encoder_length:],
      )
  
      # LSTM
      # calculate initial state
      input_hidden = (
          self.static_context_hidden_encoder_grn(static_embedding)
          .expand(self.lstm_layers, -1, -1)
          .contiguous()
      )
      input_cell = (
          self.static_context_cell_encoder_grn(static_embedding)
          .expand(self.lstm_layers, -1, -1)
          .contiguous()
      )
  
      # run local lstm encoder
      encoder_out, (hidden, cell) = self.lstm_encoder(
          input=embeddings_varying_encoder,
          hx=(input_hidden, input_cell),
      )
  
      # run local lstm decoder
      decoder_out, _ = self.lstm_decoder(
          input=embeddings_varying_decoder, hx=(hidden, cell)
      )
  
      lstm_layer = torch.cat([encoder_out, decoder_out], dim=dim_time)
      input_embeddings = torch.cat(
          [embeddings_varying_encoder, embeddings_varying_decoder],
          dim=dim_time,
      )
  
      # post lstm GateAddNorm
      lstm_out = self.post_lstm_gan(x=lstm_layer, skip=input_embeddings)
  
      # static enrichment
      static_context_enriched = self.static_context_enrichment(
          static_embedding
      )
      attn_input = self.static_enrichment_grn(
          x=lstm_out,
          context=self.expand_static_context(
              context=static_context_enriched, time_steps=time_steps
          ),
      )
  
      # multi-head attention
      attn_out, attn_out_weights = self.multihead_attn(
          q=attn_input[:, encoder_length:],
          k=attn_input,
          v=attn_input,
          mask=self.attention_mask,
      )
  
      # skip connection over attention
      attn_out = self.post_attn_gan(
          x=attn_out,
          skip=attn_input[:, encoder_length:],
      )
  
      # feed-forward
      out = self.feed_forward_block(x=attn_out)
  
      # skip connection over temporal fusion decoder from LSTM post _GateAddNorm
      out = self.pre_output_gan(
          x=out,
          skip=lstm_out[:, encoder_length:],
      )
  
      # generate output for n_targets and loss_size elements for loss evaluation
      out = self.output_layer(out)
      out = out.view(
          batch_size,
          self.output_chunk_length,
          self.n_targets,
          self.loss_size,
      )
      self._attn_out_weights = attn_out_weights
      self._static_covariate_var = static_covariate_var
      self._encoder_sparse_weights = encoder_sparse_weights
      self._decoder_sparse_weights = decoder_sparse_weights
      return out
  
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 12 for calling the _TFTModule obj's underlying member method code with example usage: obj.forward(x_in) and returns: out