code: def confusion_matrix(
      y_true: Array,
      y_pred: Array,
      num_classes: int,
      weights: Optional[Array] = None,
      np_backbone: PyModule = jnp,
  ) -> Array:
      """Computes the confusion matrix between y_true and y_pred.
  
      Args:
        y_true: Array of true labels.
        y_pred: Array of predicted labels.
        num_classes: Number of classes.
        weights: nd-array, Weight of each datapoint (e.g. for masking).
        np_backbone: numpy module: Either the regular numpy package or jax.numpy.
  
      Returns:
        A [num_classes, num_classes] confusion matrix, normalized by the number of
          elements in y_true/y_pred.
      """
      assert y_true.shape == y_pred.shape
      if weights is None:
          weights = np_backbone.ones_like(y_true)
      else:
          assert y_true.shape == weights.shape
  
      # If weights are all zero, histogram2d returns NaN. To avoid this, set weights
      # to 1 and then set output to zero below:
      weights_all_zero = 1.0 - np_backbone.any(weights).astype(
          np_backbone.float32
      )
      weights = weights + weights_all_zero
  
      cm, *_ = np_backbone.histogram2d(
          y_true.ravel(),
          y_pred.ravel(),
          bins=np_backbone.arange(num_classes + 1),
          weights=None if weights is None else weights.ravel(),
      )
  
      # If weights are all zero, set the confusion matrix to zero:
      cm = cm * (1.0 - weights_all_zero)
      return cm
  
imports: import functools
  from typing import Optional, Any, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 24 for calling a global function with example usage: confusion_matrix(y_true, y_pred, num_classes, weights, np_backbone) and returns: cm