code: def nucleaus_filter(logits, top_p=0.9, top_k=None):
      sorted_logits = jnp.sort(logits)[:, ::-1]  # sort descending
      sorted_indices = jnp.argsort(logits)[:, ::-1]
      cumulative_probs = jnp.cumsum(
          jax.nn.softmax(sorted_logits), axis=-1
      )
  
      if top_k is not None:
          # Keep only top_k tokens
          indices_range = jnp.arange(len(sorted_indices[0]))
          indices_range = jnp.stack(
              [indices_range] * len(sorted_indices), axis=0
          )
  
          sorted_indices_to_remove = jnp.where(
              indices_range >= top_k, sorted_indices, 0
          )
  
          _, indices_to_remove = jax.lax.sort_key_val(
              sorted_indices, sorted_indices_to_remove
          )
  
          logit_mask = 1e10 * indices_to_remove
  
          logits -= logit_mask
  
      # Remove tokens with cumulative probability above a threshold
      sorted_indices_to_remove = cumulative_probs > top_p
      sorted_indices_to_remove = jnp.concatenate(
          (
              jnp.zeros_like(sorted_indices_to_remove[:, :1]),
              sorted_indices_to_remove,
          ),
          axis=-1,
      )[:, :-1]
  
      _, indices_to_remove = jax.lax.sort_key_val(
          sorted_indices, sorted_indices_to_remove
      )
  
      logit_mask = 1e10 * indices_to_remove
  
      logits -= logit_mask
  
      return logits
  
imports: import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: nucleaus_filter(logits, top_p, top_k) and returns: logits