code: def weighted_recall(
      logits: Array,
      multi_hot_target: Array,
      weights: Optional[Array] = None,
  ) -> Array:
      """Computes weighted recall given the top k prediction.
  
      This computes the weighted number of correctly recalled examples/pixels in a
      single, potentially padded minibatch, given the top-k prediction. Per sample,
      k is the number of gt labels in that sample. If the minibatch/inputs is padded
      (i.e., it contains null examples/pad pixels) it is assumed that weights is a
      binary mask where 0 indicates that the example/pixel is null/padded. We assume
      the trainer will aggregate and divide by number of samples.
  
      Args:
        logits: float array; Output of model in shape [batch, ..., num_classes].
        multi_hot_target: Multi hot vector of shape [batch, ..., num_classes].
        weights: None or array of shape [batch, ...] (rank of multi_hot_target -1).
  
      Returns:
        The fraction of correctly recalled labels.
      """
      if logits.ndim != multi_hot_target.ndim:
          raise ValueError(
              "Incorrect shapes. Got shape %s logits and %s one_hot_target"
              % (str(logits.shape), str(multi_hot_target.shape))
          )
  
      num_classes = multi_hot_target.shape[-1]
  
      indices_top = jnp.argsort(logits, axis=-1)[..., ::-1]
      predictions_at_top = jax.nn.one_hot(indices_top, num_classes)
      correct_at_top = jnp.sum(
          predictions_at_top
          * jnp.expand_dims(multi_hot_target, axis=-2),
          axis=-1,
      )
  
      # Mask out (in)correct predictions that are not in top k, where k is the
      # number of gt labels.
      num_gt_labels = jnp.sum(multi_hot_target, axis=-1, keepdims=True)
      mask = (num_gt_labels > jnp.arange(num_classes)).astype(jnp.int32)
  
      recall = jnp.sum(correct_at_top * mask, axis=-1) / (
          jnp.sum(multi_hot_target, axis=-1) + 1e-12
      )
  
      if weights is not None:
          recall = apply_weights(recall, weights)
  
      return recall
  
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 8 for calling a global function with example usage: weighted_recall(logits, multi_hot_target, weights) and returns: recall