def test_reproduce_paper_evals(self, num_chunks):
    """Reproduce results from https://www.aclweb.org/anthology/P18-1009.pdf."""
    num_samples = self.labels.shape[0]
    chunk_size = num_samples // num_chunks
    metrics = []
    for chunk_start in range(0, num_samples, chunk_size):
      chunk_end = min(chunk_start + chunk_size, num_samples)
      labels = self.labels[chunk_start:chunk_end]
      predictions = self.predictions[chunk_start:chunk_end]
      current_metrics = ultra_fine_entity_typing_task.get_prediction_recall_metrics(
          labels, predictions)
      current_metrics = jax.tree.map(
          lambda x: jnp.expand_dims(x, 0), current_metrics
      )
      metrics.append(current_metrics)

    metrics = common_utils.get_metrics(metrics)
    metrics_sum = jax.tree.map(jnp.sum, metrics)
    processed_metrics = metric_utils.process_metrics(metrics_sum)
    self.assertAlmostEqual(
        processed_metrics['total_precision_value'], 0.481, places=3)
    self.assertAlmostEqual(
        processed_metrics['total_recall_value'], 0.232, places=3)
    self.assertAlmostEqual(
        processed_metrics['coarse_grained_precision_value'], 0.603, places=3)
    self.assertAlmostEqual(
        processed_metrics['coarse_grained_recall_value'], 0.616, places=3)
    self.assertAlmostEqual(
        processed_metrics['fine_grained_precision_value'], 0.404, places=3)
    self.assertAlmostEqual(
        processed_metrics['fine_grained_recall_value'], 0.384, places=3)
    self.assertAlmostEqual(
        processed_metrics['ultra_fine_grained_precision_value'],
        0.428,
        places=3)
    self.assertAlmostEqual(
        processed_metrics['ultra_fine_grained_recall_value'], 0.088, places=3)
Reproduce results from https://www.aclweb.org/anthology/P18-1009.pdf.
language/mentionmemory/tasks/ultra_fine_entity_typing_task_test.py