code: def weighted_top_one_correctly_classified(
      logits: jnp.ndarray,
      multihot_targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      background_logit_threshold: float = 0.0,
  ) -> jnp.ndarray:
      """Computes weighted number of correctly classified.
  
      We assume there are background samples where the labels are all zeros.
  
      Args:
        logits: Class logits in shape (batch_size, num_frames, num_classes).
        multihot_targets: Multihot class labels in shape (batch_size, num_frames,
          num_classes).
        weights: None or weights in shape (batch_size, num_frames).
        background_logit_threshold: If the max logit is lower than this score, we
          predict this example as background.
  
      Returns:
        Weighted numbers of correctly classified samples.
      """
      if logits.shape[-1] == 1:
          preds = logits >= 0.0
          correct = jnp.equal(preds, multihot_targets)
      else:
          top1_idx = jnp.argmax(logits, axis=-1)
          background_label = jnp.sum(multihot_targets, axis=-1) == 0
          background_pred = (
              jnp.max(logits, axis=-1) <= background_logit_threshold
          ).astype(np.int32)
  
          # Extracts the label at the highest logit index for each input.
          top1_correct = jnp.take_along_axis(
              multihot_targets, top1_idx[..., None], axis=-1
          )
          top1_correct = jnp.squeeze(top1_correct)
          foreground_correct = (
              ~background_pred.astype(bool) * top1_correct
          )
  
          # Count correctly classified background samples.
          background_correct = background_pred * background_label
          correct = foreground_correct + background_correct
  
      if weights is not None:
          return model_utils.apply_weights(correct, weights)
      return correct
  
imports: import functools
  from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Union
  import immutabledict
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.unloc import metrics
label_desc: Use image node_id 1 for calling a global function with example usage: weighted_top_one_correctly_classified(logits, multihot_targets, weights, background_logit_threshold) and returns: correct, model_utils