code: def weighted_unnormalized_box_regression_loss(
      displacements: jnp.ndarray,
      gt_displacements: jnp.ndarray,
      label: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      loss_type: str = "l1+iou",
  ) -> jnp.ndarray:
      """Computes weighted box regression losses.
  
      Box regression losses 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.
        loss_type: Box regression loss type. Multiple losses can be used
          simultaneously connected by a `+` and one can also specify the weight for
          each loss, e.g., `0.5*l1+1.0*iou`.
  
      Returns:
        The weighted DIoU losses in shape (batch_size, num_frames, num_classes) for
        TAL or (batch_size, num_captions, num_frames) for MR.
      """
  
      box_regression_loss = 0.0
      loss_types = loss_type.split("+")
      for weight_and_type in loss_types:
          wt = weight_and_type.split("*")
          (w, t) = (
              (1.0, wt[0]) if len(wt) == 1 else (float(wt[0]), wt[1])
          )
          box_regression_loss += w * _BOX_REGRESSION_LOSS_FNS[t](
              displacements, gt_displacements
          )
      # Only compute the losses on the positive segments.
      box_regression_loss = box_regression_loss * label
      if weights is None:
          return box_regression_loss
      return model_utils.apply_weights(box_regression_loss, 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 2 for calling a global function with example usage: weighted_unnormalized_box_regression_loss(displacements, gt_displacements, label, weights, loss_type) and returns: model_utils, box_regression_loss