def get_mrr(labels: Array, logits: Array) -> Array:
  """Mean reciprocal rank in https://www.aclweb.org/anthology/P18-1009.pdf."""
  labels_exists = labels.sum(axis=-1) > 0
  labels_exists = labels_exists.astype(jnp.float32)
  order = jnp.argsort(-logits, axis=-1)
  ranks = jnp.argsort(order, axis=-1)
  mrr_per_sample = 1.0 / (ranks + 1)
  mrr_per_sample = (labels * mrr_per_sample).sum(-1) / (
      labels.sum(axis=-1) + 1e-5)
  return {  # pytype: disable=bad-return-type  # jax-ndarray
      'value': jnp.dot(mrr_per_sample, labels_exists),
      'denominator': labels_exists.sum(),
  }
Mean reciprocal rank in https://www.aclweb.org/anthology/P18-1009.pdf.
language/mentionmemory/tasks/ultra_fine_entity_typing_task.py