code: def sparse_weighted_unnormalized_softmax_cross_entropy(
      logits: jnp.ndarray,
      labels: jnp.ndarray,
      mlm_weights: jnp.ndarray,
      batch_mask_weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Computes sparse weighted softmax cross entropy give logits and targets.
  
      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:
        Per example Loss value.
      """
      batch_size, length, vocab_size = logits.shape
      logits = jax.nn.log_softmax(logits)
      logits, mlm_weights = logits.ravel(), mlm_weights.ravel()
      offsets = (np.arange(batch_size * length) * vocab_size).reshape(
          (-1, length)
      )
      labels = (labels + offsets).ravel()
      loss = -jnp.take(logits, labels, axis=0)
      loss = loss * mlm_weights
      loss = loss.sum(axis=-1, keepdims=True) / (
          mlm_weights.sum(axis=-1, keepdims=True) + 1e-8
      )
      if batch_mask_weights is not None:
          loss = model_utils.apply_weights(loss, batch_mask_weights)
  
      return 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 2 for calling a global function with example usage: sparse_weighted_unnormalized_softmax_cross_entropy(logits, labels, mlm_weights, batch_mask_weights) and returns: loss