def exact_match(logits, targets, weights=None):
  """Sequence exact match averaged over batch size."""
  if weights is None:
    weights = jnp.ones(targets.shape)
  predicted_targets = logits.argmax(-1)
  match_targets = (targets == predicted_targets).astype(jnp.float32) * weights
  match_sum_per_example = match_targets.sum(1)
  expected_sum_per_example = weights.sum(1)
  example_weights = expected_sum_per_example > 0
  em = (match_sum_per_example == expected_sum_per_example) * example_weights
  em = em.sum() * 100
  normalizing_factor = example_weights.sum()
  return em, normalizing_factor
Sequence exact match averaged over batch size.
language/gscan/xattn_model/evaluation.py