def plot_heatmap(self, ax: plt.Axes | None = None, **plot_kwargs) -> plt.Axes:
        """plot the heatmap of XEB errors.

        Args:
            ax: the plt.Axes to plot on. If not given, a new figure is created,
                plotted on, and shown.
            **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.

        Returns:
            The plt.Axes that was plotted on.
        """
        show_plot = not ax
        if not isinstance(ax, plt.Axes):
            fig, ax = plt.subplots(1, 1, figsize=(8, 8))
        heatmap_data: dict[tuple[cirq.GridQubit, ...], float] = {
            pair: self.xeb_error(*pair) for pair in self.all_qubit_pairs
        }

        ax.set_title('device xeb error heatmap')

        vis.TwoQubitInteractionHeatmap(heatmap_data).plot(ax=ax, **plot_kwargs)
        if show_plot:
            fig.show()
        return ax
plot the heatmap of XEB errors.

        Args:
            ax: the plt.Axes to plot on. If not given, a new figure is created,
                plotted on, and shown.
            **plot_kwargs: Arguments to be passed to 'plt.Axes.plot'.

        Returns:
            The plt.Axes that was plotted on.
        
cirq-core/cirq/experiments/two_qubit_xeb.py