code: def verb_hard_neg_nce(
      encoded_video: jnp.ndarray,
      encoded_text: jnp.ndarray,
      mask_text: jnp.ndarray,
      temperature: float = 0.05,
      v2t_weight: float = 1.0,
      t2v_weight: float = 1.0,
      beta: float = 0.0,
  ) -> float:
      """Returns HN-NCE loss when including verb hard negatives.
  
      Args:
        encoded_video: The encoded videos.
        encoded_text: The encoded text.
        mask_text: 0 where caption (pos or neg), 1 where padding,
          shape [total_num_texts, 1]
        temperature: Temperature for scaling softmax distribution.
        v2t_weight: weight on vid2text loss.
        t2v_weight: weight on text2vid loss
        beta: beta parameter used for HN-NCE: see https://arxiv.org/abs/2301.02280
      """
      logits = utils.compute_inners(encoded_video, encoded_text)
      (
          labels,
          masking,
          inverse_mask_hn_other_vid,
      ) = get_contrastive_labels(logits, mask_text)
      logits = logits / temperature
  
      inverse_mask_hn_other_vid = jnp.sum(
          inverse_mask_hn_other_vid, axis=-1, keepdims=True
      )
      # Step 1): Video-to-text loss
      v2t_loss = hn_nce_loss_with_hardnegs(
          logits,
          labels,
          inverse_mask_hn_other_vid,
          beta=beta,
          masking=masking,
      )
      # Actually it's logits_masked.shape[0] + num neg cap per video.
      v2t_loss = v2t_loss / jnp.log(inverse_mask_hn_other_vid)
      v2t_loss = jnp.mean(v2t_loss)
  
      # Step 2): Text-to-video loss
      logits_t = jnp.transpose(logits)
      labels_t = jnp.transpose(labels)
      t2v_loss = jnp.sum(
          hn_nce_loss_with_hardnegs(logits_t, labels_t, beta=beta)
      )
      # divide by the number of pos_texts, i.e. num_vids
      t2v_loss = t2v_loss / logits.shape[0]
      # divide by log(K) where K=num_vids
      t2v_loss = t2v_loss / jnp.log(logits.shape[0])
  
      # Step 3): Add both loss terms.
      loss_final = v2t_weight * v2t_loss + t2v_weight * t2v_loss
      return loss_final
  
file_path: scenic/scenic/projects/verbs_in_action/losses.py
imports: from typing import Tuple
  import flax.linen
  import jax.numpy
  import numpy
  from scenic.projects.verbs_in_action import utils
label_desc: Use image node_id 2 for calling a global function with example usage: verb_hard_neg_nce(encoded_video, encoded_text, mask_text, temperature, v2t_weight, t2v_weight, beta) and returns: loss_final