docstr: input dimensions: n_samples x time x variables
docstr_tok: ["input","dimensions",":","n_samples","x","time","x","variables"]
code: def forward(
      self, x: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """
      input dimensions: n_samples x time x variables
      """
      encoder_lengths = x["encoder_lengths"]
      decoder_lengths = x["decoder_lengths"]
      x_cat = torch.cat(
          [x["encoder_cat"], x["decoder_cat"]], dim=1
      )  # concatenate in time dimension
      x_cont = torch.cat(
          [x["encoder_cont"], x["decoder_cont"]], dim=1
      )  # concatenate in time dimension
      timesteps = x_cont.size(1)  # encode + decode length
      max_encoder_length = int(encoder_lengths.max())
      input_vectors = self.input_embeddings(x_cat)
      input_vectors.update(
          {
              name: x_cont[..., idx].unsqueeze(-1)
              for idx, name in enumerate(self.hparams.x_reals)
              if name in self.reals
          }
      )
  
      # Embedding and variable selection
      if len(self.static_variables) > 0:
          # static embeddings will be constant over entire batch
          static_embedding = {
              name: input_vectors[name][:, 0]
              for name in self.static_variables
          }
          (
              static_embedding,
              static_variable_selection,
          ) = self.static_variable_selection(static_embedding)
      else:
          static_embedding = torch.zeros(
              (x_cont.size(0), self.hparams.hidden_size),
              dtype=self.dtype,
              device=self.device,
          )
          static_variable_selection = torch.zeros(
              (x_cont.size(0), 0), dtype=self.dtype, device=self.device
          )
  
      static_context_variable_selection = self.expand_static_context(
          self.static_context_variable_selection(static_embedding),
          timesteps,
      )
  
      embeddings_varying_encoder = {
          name: input_vectors[name][:, :max_encoder_length]
          for name in self.encoder_variables
      }
      (
          embeddings_varying_encoder,
          encoder_sparse_weights,
      ) = self.encoder_variable_selection(
          embeddings_varying_encoder,
          static_context_variable_selection[:, :max_encoder_length],
      )
  
      embeddings_varying_decoder = {
          name: input_vectors[name][:, max_encoder_length:]
          for name in self.decoder_variables  # select decoder
      }
      (
          embeddings_varying_decoder,
          decoder_sparse_weights,
      ) = self.decoder_variable_selection(
          embeddings_varying_decoder,
          static_context_variable_selection[:, max_encoder_length:],
      )
  
      # LSTM
      # calculate initial state
      input_hidden = self.static_context_initial_hidden_lstm(
          static_embedding
      ).expand(self.hparams.lstm_layers, -1, -1)
      input_cell = self.static_context_initial_cell_lstm(
          static_embedding
      ).expand(self.hparams.lstm_layers, -1, -1)
  
      # run local encoder
      encoder_output, (hidden, cell) = self.lstm_encoder(
          embeddings_varying_encoder,
          (input_hidden, input_cell),
          lengths=encoder_lengths,
          enforce_sorted=False,
      )
  
      # run local decoder
      decoder_output, _ = self.lstm_decoder(
          embeddings_varying_decoder,
          (hidden, cell),
          lengths=decoder_lengths,
          enforce_sorted=False,
      )
  
      # skip connection over lstm
      lstm_output_encoder = self.post_lstm_gate_encoder(encoder_output)
      lstm_output_encoder = self.post_lstm_add_norm_encoder(
          lstm_output_encoder, embeddings_varying_encoder
      )
  
      lstm_output_decoder = self.post_lstm_gate_decoder(decoder_output)
      lstm_output_decoder = self.post_lstm_add_norm_decoder(
          lstm_output_decoder, embeddings_varying_decoder
      )
  
      lstm_output = torch.cat(
          [lstm_output_encoder, lstm_output_decoder], dim=1
      )
  
      # static enrichment
      static_context_enrichment = self.static_context_enrichment(
          static_embedding
      )
      attn_input = self.static_enrichment(
          lstm_output,
          self.expand_static_context(
              static_context_enrichment, timesteps
          ),
      )
  
      # Attention
      attn_output, attn_output_weights = self.multihead_attn(
          q=attn_input[
              :, max_encoder_length:
          ],  # query only for predictions
          k=attn_input,
          v=attn_input,
          mask=self.get_attention_mask(
              encoder_lengths=encoder_lengths,
              decoder_lengths=decoder_lengths,
          ),
      )
  
      # skip connection over attention
      attn_output = self.post_attn_gate_norm(
          attn_output, attn_input[:, max_encoder_length:]
      )
  
      output = self.pos_wise_ff(attn_output)
  
      # skip connection over temporal fusion decoder (not LSTM decoder despite the LSTM output contains
      # a skip from the variable selection network)
      output = self.pre_output_gate_norm(
          output, lstm_output[:, max_encoder_length:]
      )
      if self.n_targets > 1:  # if to use multi-target architecture
          output = [
              output_layer(output) for output_layer in self.output_layer
          ]
      else:
          output = self.output_layer(output)
  
      return self.to_network_output(
          prediction=self.transform_output(
              output, target_scale=x["target_scale"]
          ),
          encoder_attention=attn_output_weights[
              ..., :max_encoder_length
          ],
          decoder_attention=attn_output_weights[
              ..., max_encoder_length:
          ],
          static_variables=static_variable_selection,
          encoder_variables=encoder_sparse_weights,
          decoder_variables=decoder_sparse_weights,
          decoder_lengths=decoder_lengths,
          encoder_lengths=encoder_lengths,
      )
  
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 5 for calling the TemporalFusionTransformer obj's underlying member method code with example usage: obj.forward(x) and returns: self