def plot(
        self, ax: plt.Axes | None = None, **kwargs: Any
    ) -> tuple[plt.Axes, mpl_collections.Collection]:
        """Plots the heatmap on the given Axes.
        Args:
            ax: the Axes to plot on. If not given, a new figure is created,
                plotted on, and shown.
            **kwargs: The optional keyword arguments are used to temporarily
                override the values present in the heatmap config. See
                __init__ for more details on the allowed arguments.
        Returns:
            A 2-tuple ``(ax, collection)``. ``ax`` is the `plt.Axes` that
            is plotted on. ``collection`` is the collection of paths drawn and filled.
        """
        show_plot = not ax
        if ax is None:
            fig, ax = plt.subplots(figsize=(8, 8))
        original_config = copy.deepcopy(self._config)
        self.update_config(**kwargs)

        highlighted_qubits = frozenset(kwargs.get("highlighted_qubits", ()))
        if highlighted_qubits:
            edgecolors = tuple(
                (
                    "red"
                    if not highlighted_qubits.isdisjoint(qubits)
                    else self._config["collection_options"].get("edgecolors", "grey")
                )
                for qubits in sorted(self._value_map.keys())
            )
            linestyles = tuple(
                (
                    "solid"
                    if not highlighted_qubits.isdisjoint(qubits)
                    else self._config["collection_options"].get("linestyles", "dashed")
                )
                for qubits in sorted(self._value_map.keys())
            )
            linewidths = tuple(
                (
                    4
                    if not highlighted_qubits.isdisjoint(qubits)
                    else self._config["collection_options"].get("linewidths", 2)
                )
                for qubits in sorted(self._value_map.keys())
            )
            self._config["collection_options"].update(
                {"edgecolors": edgecolors, "linestyles": linestyles, "linewidths": linewidths}
            )

        collection = self._plot_on_axis(ax)
        if show_plot:
            fig.show()
        self._config = original_config
        return (ax, collection)
Plots the heatmap on the given Axes.
        Args:
            ax: the Axes to plot on. If not given, a new figure is created,
                plotted on, and shown.
            **kwargs: The optional keyword arguments are used to temporarily
                override the values present in the heatmap config. See
                __init__ for more details on the allowed arguments.
        Returns:
            A 2-tuple ``(ax, collection)``. ``ax`` is the `plt.Axes` that
            is plotted on. ``collection`` is the collection of paths drawn and filled.
        
cirq-core/cirq/vis/heatmap.py