code: def get_confusion_matrix(
      labels: Array, logits: Array, batch_mask: Array
  ) -> Array:
      """Computes confusion matrix from predictions.
  
      Args:
        labels: [n_batch] or [n_batch, n_classes] array. In the latter case, labels
          are assumed to be one-hot, since the confusion matrix is only defined when
          each example has one label.
        logits: [n_batch, n_classes] array, which are the predictions of the model.
        batch_mask: [n_batch] array. Entries should be 1 or 0, and indicate if the
          example is valid or not.
  
      Returns:
        confusion_matrix of shape [1, n_classes, n_classes]
      """
      if labels.ndim == logits.ndim:  # one-hot targets
          y_true = jnp.argmax(labels, axis=-1)
      else:
          y_true = labels
      y_pred = jnp.argmax(logits, axis=-1)
  
      # Prepare sample weights for confusion matrix:
      weights = batch_mask.astype(jnp.float32)
  
      confusion_matrix = model_utils.confusion_matrix(
          y_true=y_true,
          y_pred=y_pred,
          num_classes=logits.shape[-1],
          weights=weights,
      )
      confusion_matrix = confusion_matrix[
          jnp.newaxis, ...
      ]  # Dummy batch dim.
      return confusion_matrix
  
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 5 for calling a global function with example usage: get_confusion_matrix(labels, logits, batch_mask) and returns: confusion_matrix