code: def sparse_weighted_per_example_accuracy(
      logits: jnp.ndarray,
      labels: jnp.ndarray,
      mlm_weights: jnp.ndarray,
      batch_mask_weights: Optional[jnp.ndarray] = None,
  ) -> jnp.ndarray:
      """Computes weighted number of correctly classified over the given batch.
  
      This computes the weighted number of correctly classified masked tokens in a
      single, potentially padded minibatch. If the minibatch/inputs is padded (i.e.,
      it contains null examples/pad pixels) it is assumed that batch_mask_weights
      is a binary mask where 0 indicates that the example/pixel is null/padded.
      We assume the trainer will aggregate and divide by number of samples.
  
      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 accuracy of predicted masked tokens.
      """
      preds = jnp.argmax(logits, axis=-1)
      correct = jnp.equal(preds, labels)
      correct = correct * mlm_weights
      # Shape of per example acccuracy will be (batch_size,).
      per_ex_accuracy = correct.sum(axis=-1) / (
          mlm_weights.sum(axis=-1) + 1e-8
      )
      if batch_mask_weights is not None:
          per_ex_accuracy = model_utils.apply_weights(
              per_ex_accuracy, batch_mask_weights
          )
      return per_ex_accuracy
  
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 4 for calling a global function with example usage: sparse_weighted_per_example_accuracy(logits, labels, mlm_weights, batch_mask_weights) and returns: per_ex_accuracy