src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py/FlaxWav2Vec2GumbelVectorQuantizer/_compute_perplexity
class FlaxWav2Vec2GumbelVectorQuantizer:    def _compute_perplexity(probs, mask=None):
        if mask is not None:
            mask_extended = jnp.broadcast_to(mask.flatten()[:, None, None], probs.shape)
            probs = jnp.where(mask_extended, probs, jnp.zeros_like(probs))
            marginal_probs = probs.sum(axis=0) / mask.sum()
        else:
            marginal_probs = probs.mean(axis=0)

        perplexity = jnp.exp(-jnp.sum(marginal_probs * jnp.log(marginal_probs + 1e-7), axis=-1)).sum()
        return perplexity