code: def weighted_top_one_correctly_classified(
      logits: jnp.ndarray,
      multi_hot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Computes weighted number of correctly classified, given top 1 class.
  
      This computes the weighted number of correctly classified examples/pixels in a
      single, potentially padded minibatch, given top-one prediction. 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: Output of model in shape [batch, ..., num_classes].
       multi_hot_targets: Multi hot vector of shape [batch, ..., num_classes].
       weights: None or array of shape [batch, ...] (rank of one_hot_targets -1).
  
      Returns:
        The number of correctly classified examples in the given batch, given top
        one prediction.
      """
      if logits.ndim != multi_hot_targets.ndim:
          raise ValueError(
              "Incorrect shapes. Got shape %s logits and %s multi_hot_targets"
              % (str(logits.shape), str(multi_hot_targets.shape))
          )
  
      top1_idx = jnp.argmax(logits, axis=-1)[..., None]
      # Extracts the label at the highest logit index for each input.
      top1_correct = jnp.take_along_axis(
          multi_hot_targets, top1_idx, axis=-1
      )
      if weights is not None:
          top1_correct = apply_weights(top1_correct, weights)
  
      return top1_correct
  
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 5 for calling a global function with example usage: weighted_top_one_correctly_classified(logits, multi_hot_targets, weights) and returns: top1_correct