def cross_entropy(logits, labels, reduction=None):
        """
        Args:
            logits: bsz, seqlen, vocab_size
            labels: bsz, seqlen
        """
        vocab_size = logits.shape[-1]
        labels = (labels[..., None] == jnp.arange(vocab_size)[None]).astype("f4")
        logits = jax.nn.log_softmax(logits, axis=-1)
        loss = -jnp.sum(labels * logits, axis=-1)
        if reduction is not None:
            loss = reduction(loss)
        return loss

        Args:
            logits: bsz, seqlen, vocab_size
            labels: bsz, seqlen
        
dpkd/transformers/examples/research_projects/jax-projects/big_bird/bigbird_flax.py