def compute_weighted_cross_entropy(
    scores: Array,
    targets: Array,
    weights: Array,
    inputs_are_prob: Optional[bool] = False,
) -> Tuple[float, float]:
  """Compute weighted cross entropy and entropy for log probs and targets.

  Args:
   scores: [batch, length, num_classes] float array.
   targets: [batch, length] categorical target integer array.
   weights: [batch, length].
   inputs_are_prob: true if inputs are probabilities rather than logits.

  Returns:
    Tuple of scalar loss and batch denominator.
  """
  scores = scores.astype(jnp.float32)
  targets = targets.astype(jnp.float32)
  weights = weights.astype(jnp.float32)
  if scores.ndim != targets.ndim + 1:
    raise ValueError('Incorrect shapes. Got shape %s scores and %s targets' %
                     (str(scores.shape), str(targets.shape)))
  vocab_size = scores.shape[-1]
  soft_targets = jax.nn.one_hot(targets, vocab_size)

  if inputs_are_prob:
    loss = -jnp.sum(soft_targets * jnp.log(scores + _SMALL_NUMBER), axis=-1)
  else:
    loss = -jnp.sum(soft_targets * jax.nn.log_softmax(scores), axis=-1)

  loss = loss * weights
  normalizing_factor = weights.sum()

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

  Args:
   scores: [batch, length, num_classes] float array.
   targets: [batch, length] categorical target integer array.
   weights: [batch, length].
   inputs_are_prob: true if inputs are probabilities rather than logits.

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