def accuracy(logits, targets, weights=None):
    """Compute weighted accuracy for log probs and targets.
    Args:
     logits: [batch, length, num_classes] float array.
     targets: categorical targets [batch, length] int array.
     weights: None or array of shape [batch, length]
    Returns:
      Tuple of scalar loss and batch normalizing factor.
    """
    if logits.ndim != targets.ndim + 1:
        raise ValueError(
            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
        )

    loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
    loss *= weights

    return loss.sum(), weights.sum()
Compute weighted accuracy for log probs and targets.
    Args:
     logits: [batch, length, num_classes] float array.
     targets: categorical targets [batch, length] int array.
     weights: None or array of shape [batch, length]
    Returns:
      Tuple of scalar loss and batch normalizing factor.
    
dpkd/transformers/examples/research_projects/performer/run_mlm_performer.py