code: def weighted_unnormalized_iou(
      displacements: jnp.ndarray,
      gt_displacements: jnp.ndarray,
      label: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Computes weighted IoUs.
  
      IoUs are only computed at frames within a positive segment.
  
      Args:
        displacements: Predicted start/end time displacements in shape (batch_size,
          num_frames, num_classes, 2) for TAL or (batch_size, num_captions,
          num_frames, 2) for MR.
        gt_displacements: Ground truth start/end time displacements in shape
          (batch_size, num_frames, num_classes, 2) for TAL or (batch_size,
          num_captions, num_frames, 2) for MR.
        label: Multihot vector of shape (batch_size, num_frames, num_classes) for
          TAL or (batch_size, num_captions, num_frames) for MR.
        weights: None or weights array of shape (batch_size, num_frames) for TAL or
          (batch_size, num_captions, num_frames) for MR.
  
      Returns:
        The weighted IoUs in a batch.
      """
  
      iou = unloc_metrics.temporal_iou(displacements, gt_displacements)
      # Only compute the losses on the positive segments.
      iou = iou * label
      return model_utils.apply_weights(iou, weights)
  
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 3 for calling a global function with example usage: weighted_unnormalized_iou(displacements, gt_displacements, label, weights) and returns: model_utils