def compute_loss_and_prob_from_probs_with_duplicates(
    probs: Array,
    classes: Array,
    targets: Array,
    weights: Array,
) -> Tuple[float, float, float]:
  """Compute weighted loss and avg correct probability given probs and targets.

  Args:
   probs: [batch, length, num_items] float array.
   classes: [batch, length, num_items] class for each item.
   targets: [batch, length] categorical target int array.
   weights:  [batch, length].

  Returns:
    Tuple of scalar loss, avg correct probability and normalizing factor.
  """
  probs = probs.astype(jnp.float32)
  weights = weights.astype(jnp.float32)

  correct_mask = (classes == jnp.expand_dims(targets, axis=-1))
  correct_mask = correct_mask.astype(jnp.float32)

  correct_probs = (correct_mask * probs).sum(axis=-1)
  avg_probs = correct_probs * weights
  loss = -jnp.log(correct_probs + _SMALL_NUMBER)
  loss = loss * weights

  return loss.sum(), avg_probs.sum(), weights.sum()  # pytype: disable=bad-return-type  # jax-ndarray
Compute weighted loss and avg correct probability given probs and targets.

  Args:
   probs: [batch, length, num_items] float array.
   classes: [batch, length, num_items] class for each item.
   targets: [batch, length] categorical target int array.
   weights:  [batch, length].

  Returns:
    Tuple of scalar loss, avg correct probability and normalizing factor.
  
language/mentionmemory/utils/metric_utils.py