docstr: Plots the attention heads of the `TFTModel`.
  
  Parameters
  ----------
  expl_result
      A `TFTExplainabilityResult` object. Corresponds to the output of :func:`explain() <TFTExplainer.explain>`.
  plot_type
      The type of attention head plot. One of ("all", "time", "heatmap").
      If "all", will plot the attention per horizon (given the horizons in the `TFTExplainabilityResult`).
      The maximum horizon corresponds to the `output_chunk_length` of the trained `TFTModel`.
      If "time", will plot the mean attention over all horizons.
      If "heatmap", will plot the attention per horizon on a heat map. The horizons are shown on the y-axis,
      and times / relative indices on the x-axis.
  show_index_as
      The type of index to be shown. One of ("relative", "time").
      If "relative", will plot the x-axis from `(-input_chunk_length, output_chunk_length - 1)`. `0` corresponds
      to the first prediction point.
      If "time", will plot the x-axis with the actual time index (or range index) of the corresponding
      `TFTExplainabilityResult`.
  ax
      Optionally, an axis to plot on. Only effective on a single `expl_result`.
  max_nr_series
      The maximum number of plots to show in case `expl_result` was computed on multiple series.
  show_plot
      Whether to show the plot.
docstr_tok: ["Plots","the","attention","heads","of","the","`","TFTModel","`",".","Parameters","--","--","--","--","--","expl_result","A","`","TFTExplainabilityResult","`","object",".","Corresponds","to","the","output","of",":","func",":","`","explain","(",")","<","TFTExplainer.explain",">","`",".","plot_type","The","type","of","attention","head","plot",".","One","of","(","``","all","''",",","``","time","''",",","``","heatmap","''",")",".","If","``","all","''",",","will","plot","the","attention","per","horizon","(","given","the","horizons","in","the","`","TFTExplainabilityResult","`",")",".","The","maximum","horizon","corresponds","to","the","`","output_chunk_length","`","of","the","trained","`","TFTModel","`",".","If","``","time","''",",","will","plot","the","mean","attention","over","all","horizons",".","If","``","heatmap","''",",","will","plot","the","attention","per","horizon","on","a","heat","map",".","The","horizons","are","shown","on","the","y-axis",",","and","times","\/","relative","indices","on","the","x-axis",".","show_index_as","The","type","of","index","to","be","shown",".","One","of","(","``","relative","''",",","``","time","''",")",".","If","``","relative","''",",","will","plot","the","x-axis","from","`","(","-input_chunk_length",",","output_chunk_length","-","1",")","`",".","`","0","`","corresponds","to","the","first","prediction","point",".","If","``","time","''",",","will","plot","the","x-axis","with","the","actual","time","index","(","or","range","index",")","of","the","corresponding","`","TFTExplainabilityResult","`",".","ax","Optionally",",","an","axis","to","plot","on",".","Only","effective","on","a","single","`","expl_result","`",".","max_nr_series","The","maximum","number","of","plots","to","show","in","case","`","expl_result","`","was","computed","on","multiple","series",".","show_plot","Whether","to","show","the","plot","."]
code: def plot_attention(
      self,
      expl_result: TFTExplainabilityResult,
      plot_type: Optional[Literal["all", "time", "heatmap"]] = "all",
      show_index_as: Literal["relative", "time"] = "relative",
      ax: Optional[matplotlib.axes.Axes] = None,
      max_nr_series: int = 5,
      show_plot: bool = True,
  ) -> matplotlib.axes.Axes:
      """Plots the attention heads of the `TFTModel`.
  
      Parameters
      ----------
      expl_result
          A `TFTExplainabilityResult` object. Corresponds to the output of :func:`explain() <TFTExplainer.explain>`.
      plot_type
          The type of attention head plot. One of ("all", "time", "heatmap").
          If "all", will plot the attention per horizon (given the horizons in the `TFTExplainabilityResult`).
          The maximum horizon corresponds to the `output_chunk_length` of the trained `TFTModel`.
          If "time", will plot the mean attention over all horizons.
          If "heatmap", will plot the attention per horizon on a heat map. The horizons are shown on the y-axis,
          and times / relative indices on the x-axis.
      show_index_as
          The type of index to be shown. One of ("relative", "time").
          If "relative", will plot the x-axis from `(-input_chunk_length, output_chunk_length - 1)`. `0` corresponds
          to the first prediction point.
          If "time", will plot the x-axis with the actual time index (or range index) of the corresponding
          `TFTExplainabilityResult`.
      ax
          Optionally, an axis to plot on. Only effective on a single `expl_result`.
      max_nr_series
          The maximum number of plots to show in case `expl_result` was computed on multiple series.
      show_plot
          Whether to show the plot.
      """
      single_series = False
      attentions = expl_result.get_explanation(component="attention")
      if isinstance(attentions, TimeSeries):
          attentions = [attentions]
          single_series = True
  
      for idx, attention in enumerate(attentions):
          if ax is None or not single_series:
              fig, ax = plt.subplots()
          if show_index_as == "relative":
              x_ticks = generate_index(
                  start=-self.model.input_chunk_length, end=self.n - 1
              )
              attention = TimeSeries.from_times_and_values(
                  times=generate_index(
                      start=-self.model.input_chunk_length,
                      end=self.n - 1,
                  ),
                  values=attention.values(copy=False),
                  columns=attention.components,
              )
              x_label = "Index relative to first prediction point"
          elif show_index_as == "time":
              x_ticks = attention.time_index
              x_label = "Time index"
          else:
              x_label, x_ticks = None, None
              raise_log(
                  ValueError(
                      "`show_index_as` must either be 'relative', or 'time'."
                  )
              )
  
          prediction_start_color = "red"
          if plot_type == "all":
              ax_title = "Attention per Horizon"
              y_label = "Attention"
              attention.plot(max_nr_components=-1, ax=ax)
          elif plot_type == "time":
              ax_title = "Mean Attention"
              y_label = "Attention"
              attention.mean(1).plot(label="Mean Attention Head", ax=ax)
          elif plot_type == "heatmap":
              ax_title = "Attention Heat Map"
              y_label = "Horizon"
  
              # generate a heat map
              x, y = np.meshgrid(x_ticks, np.arange(1, self.n + 1))
              c = ax.pcolormesh(
                  x, y, attention.values().transpose(), cmap="hot"
              )
              ax.axis([x.min(), x.max(), y.max(), y.min()])
              prediction_start_color = "lightblue"
              fig.colorbar(c, ax=ax, orientation="horizontal")
          else:
              raise raise_log(
                  ValueError(
                      "`plot_type` must be either 'all', 'time' or 'heatmap'"
                  ),
                  logger=logger,
              )
  
          # draw the prediction start point
          y_min, y_max = ax.get_ylim()
          ax.vlines(
              x=x_ticks[-self.n],
              ymin=y_min,
              ymax=y_max,
              label="prediction start",
              ls="dashed",
              lw=2,
              colors=prediction_start_color,
          )
          ax.set_xlabel(x_label)
          ax.set_ylabel(y_label)
          title_suffix = (
              "" if single_series else f": series index {idx}"
          )
          ax.set_title(ax_title + title_suffix)
          ax.legend(bbox_to_anchor=(1, 1), loc="upper left")
          if show_plot:
              plt.show()
  
          if idx + 1 == max_nr_series:
              break
      return ax
  
imports: from typing import Dict, List, Optional, Sequence, Union
  import matplotlib.axes
  import matplotlib.pyplot
  import numpy
  import pandas
  from torch import Tensor
  from darts import TimeSeries
  from darts.explainability import TFTExplainabilityResult
  from darts.explainability.explainability import _ForecastingModelExplainer
  from darts.logging import get_logger, raise_log
  from darts.models import TFTModel
  from darts.utils.timeseries_generation import generate_index
label_desc: Use image node_id 4 for calling the TFTExplainer obj's underlying member method code with example usage: obj.plot_attention(expl_result, plot_type, show_index_as, ax, max_nr_series, show_plot) and returns: ax