code: def sparse_weighted_softmax_cross_entropy(
      logits: jnp.ndarray,
      labels: jnp.ndarray,
      mlm_weights: jnp.ndarray,
      batch_mask_weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Same as weighted_unnormalized, but additionally takes a mean.
  
      Args:
        logits: Logits of shape [batch_size, length, vocab_size].
        labels: Labels from {0 ... vocab_size - 1} of shape [batch_size, length].
        mlm_weights: Weights of shape [batch_size, length], indicating masked tokens
          in masked language modeling task.
        batch_mask_weights: None or array of shape [batch,] indicating masked
          examples.
  
      Returns:
        The mean cross entropy of the examples in the given batch as a scalar.
      """
      if batch_mask_weights is not None:
          normalization = batch_mask_weights.sum()
      else:
          normalization = mlm_weights.shape[0]  # Batch size.
      sparse_unnormalized_softmax_ce = sparse_weighted_unnormalized_softmax_cross_entropy(  # pylint: disable=line-too-long
          logits, labels, mlm_weights, batch_mask_weights
      )
      return jnp.sum(sparse_unnormalized_softmax_ce) / (
          normalization + 1e-8
      )
  
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 3 for calling a global function with example usage: sparse_weighted_softmax_cross_entropy(logits, labels, mlm_weights, batch_mask_weights) and returns: unknown