def sampling(logits, rng, topk=0, topp=0, temperature=1.0):
  """top-k and nucleus sampling method.

  The nucleus sampling method is proposed in the paper `The Curious Case of
  Neural Text Degeneration (https://arxiv.org/pdf/1904.09751.pdf)`

  Args:
    logits: array: [batch_size, ..., vocab_size] logits at the current position.
    rng: JAX random state.
    topk: int: only the top-k logits will be considered to sample next token.
      If topk is zero, sample from full logits.
    topp: float: the smallest number of logits whose cumulative sum of probs
                 adds up to topp. The next token will be sampled from these
                 tokens. If zero, sample from full logits.
    temperature: float: sampling temperature factor. As it approaches zero this
                        becomes equivalent to greedy sampling. As it approaches
                        inf, this becomes uniform sampling.

  Returns:
    Array of sampled sequences: [batch_size, ...]
  """
  if topp < 0 or topk < 0:
    raise ValueError("`topp` and `topk` must be non-negative")
  if topp and topk:
    raise ValueError("At most one of `topp` or `topk` can be non-zero.")
  if topp > 0:
    logits = top_p_logits(logits, topp)
  if topk > 0:
    return top_k_sampling(logits, rng, topk, temperature)
  else:
    return jax.random.categorical(rng, logits / temperature)
top-k and nucleus sampling method.

  The nucleus sampling method is proposed in the paper `The Curious Case of
  Neural Text Degeneration (https://arxiv.org/pdf/1904.09751.pdf)`

  Args:
    logits: array: [batch_size, ..., vocab_size] logits at the current position.
    rng: JAX random state.
    topk: int: only the top-k logits will be considered to sample next token.
      If topk is zero, sample from full logits.
    topp: float: the smallest number of logits whose cumulative sum of probs
                 adds up to topp. The next token will be sampled from these
                 tokens. If zero, sample from full logits.
    temperature: float: sampling temperature factor. As it approaches zero this
                        becomes equivalent to greedy sampling. As it approaches
                        inf, this becomes uniform sampling.

  Returns:
    Array of sampled sequences: [batch_size, ...]
  
videogvt/train_lib/sampling.py