def get_positives_negatives(metric_name: str, labels: Array, predictions: Array,
                            group_start: int, group_end: int) -> MetricGroups:
  """Computes metrics over precision and recall for a specific groups."""
  tp = jnp.logical_and(labels[:, group_start:group_end] == 1,
                       predictions[:, group_start:group_end] == 1).sum(-1)
  fp = jnp.logical_and(labels[:, group_start:group_end] == 0,
                       predictions[:, group_start:group_end] == 1).sum(-1)
  fn = jnp.logical_and(labels[:, group_start:group_end] == 1,
                       predictions[:, group_start:group_end] == 0).sum(-1)
  precision = tp / (tp + fp + _SMALL_NUMBER)
  precision_weight = (tp + fp) > 0
  recall = tp / (tp + fn + _SMALL_NUMBER)
  recall_weight = (tp + fn) > 0
  return {
      metric_name + '_precision': {
          'value': jnp.dot(precision, precision_weight),
          'denominator': precision_weight.sum(),
      },
      metric_name + '_recall': {
          'value': jnp.dot(recall, recall_weight),
          'denominator': recall_weight.sum(),
      }
  }
Computes metrics over precision and recall for a specific groups.
language/mentionmemory/tasks/ultra_fine_entity_typing_task.py