docstr: Returns causal mask to apply for self-attention layer.
docstr_tok: ["Returns","causal","mask","to","apply","for","self-attention","layer","."]
code: def get_attention_mask(
      self,
      encoder_lengths: torch.LongTensor,
      decoder_lengths: torch.LongTensor,
  ):
      """
      Returns causal mask to apply for self-attention layer.
      """
      decoder_length = decoder_lengths.max()
      if self.hparams.causal_attention:
          # indices to which is attended
          attend_step = torch.arange(decoder_length, device=self.device)
          # indices for which is predicted
          predict_step = torch.arange(
              0, decoder_length, device=self.device
          )[:, None]
          # do not attend to steps to self or after prediction
          decoder_mask = (
              (attend_step >= predict_step)
              .unsqueeze(0)
              .expand(encoder_lengths.size(0), -1, -1)
          )
      else:
          # there is value in attending to future forecasts if they are made with knowledge currently
          #   available
          #   one possibility is here to use a second attention layer for future attention (assuming different effects
          #   matter in the future than the past)
          #   or alternatively using the same layer but allowing forward attention - i.e. only
          #   masking out non-available data and self
          decoder_mask = (
              create_mask(decoder_length, decoder_lengths)
              .unsqueeze(1)
              .expand(-1, decoder_length, -1)
          )
      # do not attend to steps where data is padded
      encoder_mask = (
          create_mask(encoder_lengths.max(), encoder_lengths)
          .unsqueeze(1)
          .expand(-1, decoder_length, -1)
      )
      # combine masks along attended time - first encoder and then decoder
      mask = torch.cat(
          (
              encoder_mask,
              decoder_mask,
          ),
          dim=2,
      )
      return mask
  
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 4 for calling the TemporalFusionTransformer obj's underlying member method code with example usage: obj.get_attention_mask(encoder_lengths, decoder_lengths) and returns: mask