def test_loss_and_metrics_as_expected(self,
                                        seed,
                                        scale,
                                        local_n_mentions,
                                        global_n_mentions,
                                        max_num_positives=None,
                                        max_num_negatives=None):
    """Test whether loss and metrics computation produces expected values."""
    np.random.seed(seed)
    max_num_negatives = max_num_negatives or global_n_mentions
    max_num_positives = max_num_positives or global_n_mentions

    shape = (local_n_mentions, global_n_mentions)
    scores = np.random.random(shape) * scale

    num_positives = np.random.randint(
        max_num_positives + 1, size=(local_n_mentions))
    num_positives[0] = 0
    num_positives[-1] = global_n_mentions

    num_negatives = np.random.randint(
        max_num_negatives + 1, size=(local_n_mentions))
    num_negatives = np.minimum(num_negatives, global_n_mentions - num_positives)

    positives = np.zeros(shape, dtype=np.bool_)
    negatives = np.zeros(shape, dtype=np.bool_)

    for index in range(local_n_mentions):
      ids = np.random.choice(
          global_n_mentions,
          num_positives[index] + num_negatives[index],
          replace=False)
      positives[index, ids[:num_positives[index]]] = True
      negatives[index, ids[num_positives[index]:]] = True
    self.assertEqual(np.logical_and(positives, negatives).sum(), 0)

    weights = np.logical_and(num_positives > 0, num_negatives > 0)

    (actual_loss, actual_metrics, (actual_acc_per_sample,
                                   actual_weights_per_sample)
    ) = metric_utils.compute_cross_entropy_loss_with_positives_and_negatives_masks(
        scores, positives, negatives)

    expected_loss, expected_acc, expected_denom = 0, 0, 0
    expected_acc_per_sample = []
    # Consider every sample independently
    for i in range(local_n_mentions):
      if not weights[i]:
        expected_acc_per_sample.append(0)
        continue
      # Collect positive and negative scores
      positive_scores, negative_scores = [], []
      for j in range(global_n_mentions):
        if positives[i, j]:
          positive_scores.append(scores[i, j])
        if negatives[i, j]:
          negative_scores.append(scores[i, j])
      self.assertNotEmpty(positive_scores)
      self.assertNotEmpty(negative_scores)
      n_pos = len(positive_scores)
      max_negative_scores = max(negative_scores)
      current_loss, current_acc = 0, 0
      # Consider positive class per sample independently
      # and compute loss using a naive softmax op
      for pos_index in range(n_pos):
        current_scores = np.array([positive_scores[pos_index]] +
                                  negative_scores)
        current_scores = jax.nn.log_softmax(current_scores)
        current_loss += -current_scores[0]
        current_acc += int(positive_scores[pos_index] > max_negative_scores)
      expected_loss += current_loss / n_pos
      expected_acc += current_acc / n_pos
      expected_denom += 1
      expected_acc_per_sample.append(current_acc / n_pos)

    self.assertAlmostEqual(actual_loss, expected_loss, delta=1)
    self.assertAlmostEqual(actual_metrics['loss'], expected_loss, delta=1)
    self.assertAlmostEqual(actual_metrics['acc'], expected_acc, places=4)
    self.assertAlmostEqual(
        actual_metrics['denominator'], expected_denom, places=4)
    self.assertTrue(np.all(weights == actual_weights_per_sample))
    self.assertSequenceAlmostEqual(
        actual_acc_per_sample, expected_acc_per_sample, places=4)
Test whether loss and metrics computation produces expected values.
language/mentionmemory/utils/metric_utils_test.py