code: def render_confusion_matrices(
      confusion_matrices: List[Array],
      normalization_method: str = "cols",
      figsize: Tuple[int, int] = (12, 12),
      dpi: int = 100,
      font_scale: int = 3,
  ) -> Array:
      """Render confusion matrix so that it can be logged to Tensorboard.
  
      Args:
        confusion_matrices: List of [n_batch, n_class, n_class] confusion matrices.
          The first two dimensions will be summed over to get an [n_class, n_class]
          matrix for rendering.
        normalization_method: Method of normalizing the confusion matrix before
          plotting. Supported values are one of "cols", "rows" and "none".
          If any other value, no normalization is performed.
        figsize: The figure size used by matplotlib and seaborn.
        dpi: The dpi used by matplotlib and seaborn.
        font_scale: The font scale used by seaborn.
  
      Returns:
        image: Rendered image of the confusion matrix for plotting. Data type is
          uint8 and values are in range [0, 255]. Shape is
          [1, figsize * dpi, figsize * dpi, 3]
      """
      conf_matrix = np.sum(
          confusion_matrices, axis=0
      )  # Sum over eval batches.
      if conf_matrix.ndim != 3:
          raise AssertionError(
              "Expecting confusion matrix to have shape "
              f"[batch_size, num_classes, num_classes], got {conf_matrix.shape}."
          )
      conf_matrix = np.sum(
          conf_matrix, axis=0
      )  # Sum over batch dimension.
  
      if normalization_method not in {"rows", "cols", "none"}:
          logging.warning(
              "Normalizer must be one of {rows, cols, none}."
              "Defaulting to none."
          )
  
      sns.set(font_scale=font_scale)
      fig = plt.figure(figsize=figsize, dpi=dpi)
  
      # Normalize entries of the confusion matrix.
      if normalization_method == "rows":
          normalizer = conf_matrix.sum(axis=1)[:, np.newaxis]
      elif normalization_method == "cols":
          normalizer = conf_matrix.sum(axis=0)[np.newaxis, :]
      else:
          normalizer = 1
      normalized_matrix = np.nan_to_num(conf_matrix / normalizer)
  
      if np.sum(normalized_matrix) > 0:
          sns.heatmap(
              normalized_matrix,
              annot=True,
              linewidths=0.5,
              square=True,
              cbar=False,
              cmap="jet",
              annot_kws={"size": 18},
          )
          fig.tight_layout(pad=0.0)
  
      fig.canvas.draw()
      ncols, nrows = fig.canvas.get_width_height()
      image = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
      image = image.reshape(nrows, ncols, 3)
      return np.expand_dims(image, axis=0)
  
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple, Union
  from absl import logging
  from flax import jax_utils
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import jax.profiler
  import matplotlib.pyplot
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
  import seaborn
label_desc: Use image node_id 6 for calling a global function with example usage: render_confusion_matrices(confusion_matrices, normalization_method, figsize, dpi, font_scale) and returns: np