def compute_tp_fp_fn_weighted(
    predictions: Array, labels: Array, weights: Array,
    ignore_class: Optional[int]) -> Tuple[float, float, float]:
  """Compute true positives, false positives and false negatives.

  Args:
   predictions: [batch, length] categorical predictions int array.
   labels: [batch, length] categorical labels int array.
   weights: [batch, length].
   ignore_class: which class to ignore in the computations

  Returns:
    Tuple with numbers of true positive, false positive and false negative
      predictions.
  """
  true_positives = (predictions == labels)
  false_positives = jnp.logical_not(true_positives)
  false_negatives = false_positives

  if ignore_class is not None:
    dont_ignore_predictions = (predictions != ignore_class)
    dont_ignore_labels = (labels != ignore_class)
    true_positives = jnp.logical_and(true_positives, dont_ignore_predictions)
    # Exactly the same as
    # true_positives = jnp.logical_and(true_positives, dont_ignore_labels)
    # since for true positives `dont_ignore_predictions` = `dont_ignore_labels`.
    false_positives = jnp.logical_and(false_positives, dont_ignore_predictions)
    false_negatives = jnp.logical_and(false_negatives, dont_ignore_labels)

  def get_weighted_sum(values):
    values = values.astype(weights.dtype)
    return jnp.dot(values, weights)

  n_true_positive = get_weighted_sum(true_positives)
  n_false_positive = get_weighted_sum(false_positives)
  n_false_negative = get_weighted_sum(false_negatives)

  return n_true_positive, n_false_positive, n_false_negative  # pytype: disable=bad-return-type  # jnp-type
Compute true positives, false positives and false negatives.

  Args:
   predictions: [batch, length] categorical predictions int array.
   labels: [batch, length] categorical labels int array.
   weights: [batch, length].
   ignore_class: which class to ignore in the computations

  Returns:
    Tuple with numbers of true positive, false positive and false negative
      predictions.
  
language/mentionmemory/utils/metric_utils.py