code: def _cls_loss_weights(
      batch_mask: jnp.ndarray,
      caption_mask: jnp.ndarray,
      frame_mask: jnp.ndarray,
      all_gather_loss: bool,
  ) -> jnp.ndarray:
      """This function is to compute the weights for classification."""
      if all_gather_loss:
          return (
              batch_mask[:, jnp.newaxis, jnp.newaxis, jnp.newaxis]
              * batch_mask[jnp.newaxis, :, jnp.newaxis, jnp.newaxis]
              * caption_mask[jnp.newaxis, :, :, jnp.newaxis]
              * frame_mask[jnp.newaxis, :, jnp.newaxis, :].astype(
                  jnp.float32
              )
          )
      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 3 for calling a global function with example usage: _cls_loss_weights(batch_mask, caption_mask, frame_mask, all_gather_loss) and returns: unknown, unknown