@measurement.evaluate(params=0, batch=1)
def evaluate(params, batch):
    tokens = np.asarray(batch, dtype=np.uint32)
    logits = model.apply({"params": params}, jnp.asarray(tokens))
    return (np.asarray(logits.argmax(axis=-1)) == (tokens.sum(axis=1) % 2)).astype(np.float32)

@measurement.aggregate
def aggregate(outputs):
    return float(np.concatenate([np.asarray(o) for o in outputs]).mean())