code: def _box_loss_weights(
      batch_mask: jnp.ndarray,
      caption_mask: jnp.ndarray,
      frame_mask: jnp.ndarray,
  ) -> jnp.ndarray:
      """This function is to compute the weights for box loss."""
      return (
          batch_mask[:, jnp.newaxis, jnp.newaxis]
          * caption_mask[..., jnp.newaxis]
          * frame_mask[:, jnp.newaxis, :].astype(jnp.float32)
      )
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, Union
  import immutabledict
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.unloc import temporal_localization_base_model
label_desc: Use image node_id 4 for calling a global function with example usage: _box_loss_weights(batch_mask, caption_mask, frame_mask) and returns: unknown