code: def get_confusion_matrix(*, labels, logits, batch_mask):
      """Computes the confusion matrix that is necessary for global mIoU."""
      if labels.ndim == logits.ndim:  # One-hot targets.
          y_true = jnp.argmax(labels, axis=-1)
      else:
          y_true = labels
          # Set excluded pixels (label -1) to zero, because the confusion matrix
          # computation cannot deal with negative labels. They will be ignored due to
          # the batch_mask anyway:
          y_true = jnp.maximum(y_true, 0)
      y_pred = jnp.argmax(logits, axis=-1)
  
      # Prepare sample weights for confusion matrix:
      weights = batch_mask.astype(jnp.float32)
      # Normalize weights by number of samples to avoid having very large numbers in
      # the confusion matrix, which could lead to imprecise results (note that we
      # should not normalize by sum(weights) because that might differ between
      # devices/hosts):
      weights = weights / weights.size
  
      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 Any, Callable, Dict, Optional, Tuple, Type
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  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 ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.vivit import evaluation_lib
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 3 for calling a global function with example usage: get_confusion_matrix() and returns: confusion_matrix