def compute_weighted_accuracy(scores: Array, targets: Array,
                              weights: Array) -> Tuple[float, float]:
  """Compute weighted accuracy for log probs and targets.

  Args:
   scores: [batch, length, num_classes] float array.
   targets: [batch, length] categorical targets int array.
   weights: [batch, length].

  Returns:
    Tuple of scalar loss and batch normalizing factor.
  """
  if scores.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s scores and %s targets' %
                     (str(scores.shape), str(targets.shape)))
  acc = jnp.equal(jnp.argmax(scores, axis=-1), targets)
  acc = acc * weights
  normalizing_factor = weights.sum()

  return acc.sum(), normalizing_factor  # pytype: disable=bad-return-type  # jax-ndarray
Compute weighted accuracy for log probs and targets.

  Args:
   scores: [batch, length, num_classes] float array.
   targets: [batch, length] categorical targets int array.
   weights: [batch, length].

  Returns:
    Tuple of scalar loss and batch normalizing factor.
  
language/mentionmemory/utils/metric_utils.py