def accuracy(logits, targets, weights=None):
  """Sequence accuracy averaged over sequence length."""
  if logits.ndim != targets.ndim + 1:
    raise ValueError("Incorrect shapes. Got shape %s logits and %s targets" %
                     (str(logits.shape), str(targets.shape)))
  acc = jnp.equal(jnp.argmax(logits, axis=-1), targets)
  normalizing_factor = np.prod(logits.shape[:-1])
  if weights is not None:
    acc = acc * weights * 100
    normalizing_factor = weights.sum()
  return acc.sum(), normalizing_factor
Sequence accuracy averaged over sequence length.
language/gscan/xattn_model/evaluation.py