code: def weighted_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 loss.
  
      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 mean box regression loss of the examples in a batch as a scalar.
      """
  
      if weights is not None:
          normalization = model_utils.apply_weights(
              label, weights
          ).sum()
      else:
          normalization = np.prod(label.shape[:-1])
      box_loss = weighted_unnormalized_box_regression_loss(
          displacements, gt_displacements, label, weights, loss_type
      )
      return jnp.sum(box_loss) / (normalization + 1e-8)
  
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 5 for calling a global function with example usage: weighted_box_regression_loss(displacements, gt_displacements, label, weights, loss_type) and returns: unknown