def top_k_sampling(logits, rng, topk=0, temperature=1.0):
  """Only top k probable tokens should be considered for sampling."""
  topk_logits, topk_idxs = lax.top_k(logits, topk)
  topk_token = jnp.expand_dims(
      random.categorical(rng, topk_logits/temperature).astype(jnp.int32),
      axis=-1)
  # Returns the original indices corresponding to the sampled top-k tokens.
  # [batch_size, ...]
  sampled_tokens = jnp.squeeze(
      jnp.take_along_axis(topk_idxs, topk_token, axis=-1),
      axis=-1).astype(jnp.int32)
  return sampled_tokens
Only top k probable tokens should be considered for sampling.
videogvt/train_lib/sampling.py