code: def compute_bert_loss(
      mlm_logits: jnp.ndarray, nsp_logits: jnp.ndarray, batch: Batch
  ) -> float:
      """Computes BERT loss.
  
      Args:
       mlm_logits: Masked LM logits of shape [batch, length, vocab_size].
       nsp_logits: Next sentence prediction logits of shape [batch, 2].
       batch: Batch of data that has 'masked_lm_ids', 'masked_lm_weights' and
         'next_sentence_labels'.
  
      Returns:
        Loss value.
      """
      next_sentence_labels = common_utils.onehot(
          batch["next_sentence_labels"], 2
      )
      batch_weights = batch.get(
          "batch_mask"
      )  # batch_mask might not be defined
      nsp_loss = model_utils.weighted_softmax_cross_entropy(
          nsp_logits, next_sentence_labels, batch_weights
      )
      mlm_loss = sparse_weighted_softmax_cross_entropy(
          mlm_logits,
          batch["masked_lm_ids"],
          batch["masked_lm_weights"],
          batch_weights,
      )
      return nsp_loss + mlm_loss
  
file_path: scenic/scenic/projects/baselines/bert/bert_base_model.py
imports: from typing import Callable, Dict, Optional, Tuple, Union
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 6 for calling a global function with example usage: compute_bert_loss(mlm_logits, nsp_logits, batch) and returns: unknown