def top_p_logits(logits, topp=0):
  """Finds the top logits with cumulative probability >= top and mask others."""
  logits_sorted = jnp.sort(logits, axis=-1)[..., ::-1]  # sort descending
  sorted_cum_probs = jnp.cumsum(
      jax.nn.softmax(logits_sorted, axis=-1), axis=-1)  # get cumsum probs
  cutoff_index = jnp.sum(
      sorted_cum_probs < topp, axis=-1, keepdims=True)  # find cutoff index
  cutoff_logit = jnp.take_along_axis(logits_sorted, cutoff_index, axis=-1)
  logits = jnp.where(logits < cutoff_logit,
                     jnp.full_like(logits, NEG_INF), logits)
  return logits
Finds the top logits with cumulative probability >= top and mask others.
videogvt/train_lib/sampling.py