code: def heatmap(
      X,
      cmap_type="seismic",
      reduce_op="sum",
      reduce_axis=-1,
      alpha_cmap=False,
      **kwargs,
  ):
      """Creates a heatmap/color map.
  
      Create a heatmap or colormap out of the input tensor.
  
      :param X: A image tensor with 4 axes.
      :param cmap_type: The color map to use. Default 'seismic'.
      :param reduce_op: Operation to reduce the color axis.
        Either 'sum' or 'absmax'.
      :param reduce_axis: Axis to reduce.
      :param alpha_cmap: Should the alpha component of the cmap be included.
      :param kwargs: Arguments passed on to :func:`project`
      :return: The tensor as color-map.
      """
      cmap = mpl.colormaps.get_cmap(cmap_type)
  
      tmp = X
      shape = tmp.shape
  
      if reduce_op == "sum":
          tmp = tmp.sum(axis=reduce_axis)
      elif reduce_op == "absmax":
          pos_max = tmp.max(axis=reduce_axis)
          neg_max = (-tmp).max(axis=reduce_axis)
          abs_neg_max = -neg_max
          tmp = np.select(
              [pos_max >= abs_neg_max, pos_max < abs_neg_max],
              [pos_max, neg_max],
          )
      else:
          raise NotImplementedError()
  
      tmp = project(tmp, output_range=(0, 255), **kwargs).astype(
          np.int64
      )
  
      if alpha_cmap:
          tmp = cmap(tmp.flatten()).T
      else:
          tmp = cmap(tmp.flatten())[:, :3].T
      tmp = tmp.T
  
      shape = list(shape)
      shape[reduce_axis] = 3 + alpha_cmap
      return tmp.reshape(shape).astype(np.float32)
  
file_path: innvestigate/src/innvestigate/utils/visualizations.py
imports: from __future__ import annotations
  import matplotlib
  import matplotlib.pyplot
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: heatmap(X, cmap_type, reduce_op, reduce_axis, alpha_cmap) and returns: tmp