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
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.
        
darts/explainability/tft_explainer.py