def test_logit_values_as_expected(self):
    """Test whether metrics computations produce expected values."""

    logits = np.random.rand(self.batch_size, self.seq_len, self.vocab_size)
    targets = np.random.randint(
        self.vocab_size, size=(self.batch_size, self.seq_len))
    dense_targets = jax.nn.one_hot(targets, self.vocab_size)
    weights = np.random.randint(2, size=(self.batch_size, self.seq_len))

    # Check loss and denominator make sense for random values
    loss, denominator = metric_utils.compute_weighted_cross_entropy(
        logits,
        targets,
        weights,
    )

    expected_loss = -jax.nn.log_softmax(logits, axis=-1) * dense_targets
    expected_loss = (expected_loss * np.expand_dims(weights, axis=-1)).sum()
    self.assertAlmostEqual(loss, expected_loss, 1)
    self.assertAlmostEqual(denominator, weights.sum(), 1)

    # Check loss makes sense for uniform and degenerate scores
    logits = np.ones(shape=(self.batch_size, self.seq_len, self.vocab_size))
    loss, denominator = metric_utils.compute_weighted_cross_entropy(
        logits,
        targets,
        weights,
    )
    expected_loss = np.log(self.vocab_size)
    self.assertAlmostEqual(loss / denominator, expected_loss, 4)

    logits = np.zeros(shape=(self.batch_size, self.seq_len, self.vocab_size))
    logits = logits + (
        _LARGE_NUMBER * dense_targets - _LARGE_NUMBER * (1 - dense_targets))

    loss, denominator = metric_utils.compute_weighted_cross_entropy(
        logits,
        targets,
        weights,
    )
    self.assertAlmostEqual(loss / denominator, 0.0, 4)
Test whether metrics computations produce expected values.
language/mentionmemory/utils/metric_utils_test.py