code: def _colorize(
      agg,
      color_key,
      how,
      alpha,
      span,
      min_alpha,
      name,
      color_baseline,
      rescale_discrete_levels,
  ):
      if cupy and isinstance(agg.data, cupy.ndarray):
          array = cupy.array
      else:
          array = np.array
  
      if not agg.ndim == 3:
          raise ValueError("agg must be 3D")
  
      cats = agg.indexes[agg.dims[-1]]
      if not len(
          cats
      ):  # No categories and therefore no data; return an empty image
          return Image(
              np.zeros(agg.shape[0:2], dtype=np.uint32),
              dims=agg.dims[:-1],
              coords=dict(
                  [
                      (agg.dims[1], agg.coords[agg.dims[1]]),
                      (agg.dims[0], agg.coords[agg.dims[0]]),
                  ]
              ),
              name=name,
          )
  
      if color_key is None:
          raise ValueError(
              "Color key must be provided, with at least as many "
              + "colors as there are categorical fields"
          )
      if not isinstance(color_key, dict):
          color_key = dict(zip(cats, color_key))
      if len(color_key) < len(cats):
          raise ValueError(
              f"Insufficient colors provided ({len(color_key)}) for the categorical "
              f"fields available ({len(cats)})"
          )
  
      colors = [rgb(color_key[c]) for c in cats]
      rs, gs, bs = map(array, zip(*colors))
  
      # Reorient array (transposing the category dimension first)
      agg_t = agg.transpose(*((agg.dims[-1],) + agg.dims[:2]))
      data = agg_t.data.transpose([1, 2, 0])
      if isinstance(data, da.Array):
          data = data.compute()
      color_data = data.copy()
  
      # subtract color_baseline if needed
      with warnings.catch_warnings():
          warnings.filterwarnings(
              "ignore", r"All-NaN slice encountered"
          )
          baseline = (
              np.nanmin(color_data)
              if color_baseline is None
              else color_baseline
          )
      with np.errstate(invalid="ignore"):
          if baseline > 0:
              color_data -= baseline
          elif baseline < 0:
              color_data += -baseline
          if (
              color_data.dtype.kind != "u"
              and color_baseline is not None
          ):
              color_data[color_data < 0] = 0
  
      color_total = nansum_missing(color_data, axis=2)
      # dot does not handle nans, so replace with zeros
      color_data[np.isnan(data)] = 0
  
      # zero-count pixels will be 0/0, but it's safe to ignore that when dividing
      with np.errstate(divide="ignore", invalid="ignore"):
          r = (color_data.dot(rs) / color_total).astype(np.uint8)
          g = (color_data.dot(gs) / color_total).astype(np.uint8)
          b = (color_data.dot(bs) / color_total).astype(np.uint8)
  
      # special case -- to give an appropriate color when min_alpha != 0 and data=0,
      # take avg color of all non-nan categories
      color_mask = ~np.isnan(data)
      cmask_sum = np.sum(color_mask, axis=2)
  
      with np.errstate(divide="ignore", invalid="ignore"):
          r2 = (color_mask.dot(rs) / cmask_sum).astype(np.uint8)
          g2 = (color_mask.dot(gs) / cmask_sum).astype(np.uint8)
          b2 = (color_mask.dot(bs) / cmask_sum).astype(np.uint8)
  
      missing_colors = np.sum(color_data, axis=2) == 0
      r = np.where(missing_colors, r2, r)
      g = np.where(missing_colors, g2, g)
      b = np.where(missing_colors, b2, b)
  
      total = nansum_missing(data, axis=2)
      mask = np.isnan(total)
      a = _interpolate_alpha(
          data,
          total,
          mask,
          how,
          alpha,
          span,
          min_alpha,
          rescale_discrete_levels,
      )
  
      values = np.dstack([r, g, b, a]).view(np.uint32).reshape(a.shape)
      if cupy and isinstance(values, cupy.ndarray):
          # Convert cupy array to numpy for final image
          values = cupy.asnumpy(values)
  
      return Image(
          values,
          dims=agg.dims[:-1],
          coords=dict(
              [
                  (agg.dims[1], agg.coords[agg.dims[1]]),
                  (agg.dims[0], agg.coords[agg.dims[0]]),
              ]
          ),
          name=name,
      )
  
imports: from __future__ import annotations
  from collections.abc import Iterator
  from io import BytesIO
  import warnings
  import numpy
  import numba
  import toolz
  import xarray
  import dask.array
  from PIL.Image import fromarray
  from datashader.colors import rgb, Sets1to3
  from datashader.utils import nansum_missing, ngjit
label_desc: Use image node_id 6 for calling a global function with example usage: _colorize(agg, color_key, how, alpha, span, min_alpha, name, color_baseline, rescale_discrete_levels) and returns: Image, Image