code: def heatmap(
      data: np.ndarray,
      xticklabels: List[str],
      yticklabels: List[str],
      vmin: Optional[float] = None,
      vmax: Optional[float] = None,
      cmap: Union[str, matplotlib.colors.Colormap] = "magma",
      robust: Optional[bool] = False,
      annot: Optional[bool] = True,
      linewidths: float = 3,
      linecolor: str = "w",
      cbar: bool = True,
      cbar_label: str = "",
      cbar_ax: Optional["plt.Axes"] = None,
      cbar_kws: Optional[dict] = None,
      fmt: Union[str, matplotlib.ticker.Formatter] = "{x:.2f}",
      textcolors: Tuple[str, str] = ("white", "black"),
      threshold: Optional[float] = None,
      text_kws: Optional[dict] = None,
      ax: Optional["plt.Axes"] = None,
      **kwargs,
  ) -> "plt.Axes":
      """
      Constructs a heatmap with annotation.
  
      Parameters
      ----------
      data
          A 2D `numpy` array of shape `M x N`.
      yticklabels
          A list or array of length `M` with the labels for the rows.
      xticklabels
          A list or array of length `N` with the labels for the columns.
      vmin, vmax
         When using scalar data and no explicit norm, `vmin` and `vmax` define the data range that the colormap covers.
         By default, the colormap covers the complete value range of the supplied data. It is an error to use
         `vmin/vmax` when norm is given. When using RGB(A) data, parameters `vmin/vmax` are ignored.
      cmap
          The Colormap instance or registered colormap name used to map scalar data to colors. This parameter is
          ignored for RGB(A) data.
      robust
          If ``True`` and `vmin` or `vmax` are absent, the colormap range is computed with robust quantiles
          instead of the extreme values. Uses `numpy.nanpercentile`_ with `q` values set to 2 and 98, respectively.
  
          .. _numpy.nanpercentile:
              https://numpy.org/doc/stable/reference/generated/numpy.nanpercentile.html
  
      annot
          Boolean flag whether to annotate the heatmap. Default ``True``.
      linewidths
          Width of the lines that will divide each cell. Default 3.
      linecolor
          Color of the lines that will divide each cell. Default ``"w"``.
      cbar
          Boolean flag whether to draw a colorbar.
      cbar_label
          Optional label for the colorbar.
      cbar_ax
          Optional axes in which to draw the colorbar, otherwise take space from the main axes.
      cbar_kws
          An optional dictionary with arguments to `matplotlib.figure.Figure.colorbar`_.
  
          .. _matplotlib.figure.Figure.colorbar:
              https://matplotlib.org/stable/api/figure_api.html#matplotlib.figure.Figure.colorbar
  
      fmt
          Format of the annotations inside the heatmap. This should either use the string format method,
          e.g. ``"{x:.2f}"``, or be a `matplotlib.ticker.Formatter`_. Default ``"{x:.2f}"``.
  
          .. _matplotlib.ticker.Formatter:
              https://matplotlib.org/stable/api/ticker_api.html#matplotlib.ticker.Formatter
  
      textcolors
          A tuple of `matplotlib` colors. The first is used for values below a threshold,
          the second for those above. Default ``("black", "white")``.
      threshold
          Optional value in data units according to which the colors from textcolors are
          applied. If ``None`` (the default) uses the middle of the colormap as
          separation.
      text_kws
          An optional dictionary with arguments to `matplotlib.axes.Axes.text`_.
  
          .. _matplotlib.axes.Axes.text:
              https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.text.html
  
      ax
          Axes in which to draw the plot, otherwise use the currently-active axes.
      kwargs
          All other keyword arguments are passed to `matplotlib.axes.Axes.imshow`_.
  
          .. _matplotlib.axes.Axes.imshow:
              https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.imshow.html
  
      Returns
      -------
      Axes object with the heatmap.
      """
  
      # get heatmap min value
      if vmin is None:
          vmin = (
              np.nanpercentile(data, 2) if robust else np.nanmin(data)
          )
  
      # get heatmap max value
      if vmax is None:
          vmax = (
              np.nanpercentile(data, 98) if robust else np.nanmax(data)
          )
  
      # create the heatmap
      ax = _create_heatmap(
          data=data,
          yticklabels=yticklabels,
          xticklabels=xticklabels,
          ax=ax,
          vmin=vmin,
          vmax=vmax,
          cmap=cmap,
          linewidths=linewidths,
          linecolor=linecolor,
          cbar=cbar,
          cbar_kws=cbar_kws,
          cbar_ax=cbar_ax,
          cbar_label=cbar_label,
          **kwargs,
      )
  
      if annot:
          if text_kws is None:
              text_kws = {}
  
          # annotate the heatmap
          _annotate_heatmap(
              im=ax.get_images()[0],
              fmt=fmt,
              textcolors=textcolors,
              threshold=threshold,
              **text_kws,
          )
      return ax
  
imports: import warnings
  from enum import Enum
  from typing import List, Optional, Tuple, Union
  import matplotlib
  import numpy
  from matplotlib import pyplot
  from matplotlib.colors import LinearSegmentedColormap
  from matplotlib.pyplot import Axes, Figure
  from mpl_toolkits.axes_grid1 import make_axes_locatable
  from numpy import ndarray
label_desc: Use image node_id 8 for calling a global function with example usage: heatmap(data, xticklabels, yticklabels, vmin, vmax, cmap, robust, annot, linewidths, linecolor, cbar, cbar_label, cbar_ax, cbar_kws, fmt, textcolors, threshold, text_kws, ax) and returns: ax