class_docstr: [`FlaxLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off.
  
  Args:
      top_p (`float`):
          If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
          higher are kept for generation.
      filter_value (`float`, *optional*, defaults to -inf):
          All filtered values will be set to this float value.
      min_tokens_to_keep (`int`, *optional*, defaults to 1):
          Minimum number of tokens that cannot be filtered.
code: def __call__(
      self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int
  ) -> jnp.ndarray:
      topk_scores, topk_indices = lax.top_k(scores, scores.shape[-1])
  
      mask_scores = jnp.full_like(scores, self.filter_value)
      cumulative_probs = jax.nn.softmax(topk_scores, axis=-1).cumsum(
          axis=-1
      )
      score_mask = cumulative_probs < self.top_p
  
      # include the token that is higher than top_p as well
      score_mask = jnp.roll(score_mask, 1)
      score_mask |= score_mask.at[:, 0].set(True)
  
      # min tokens to keep
      score_mask = score_mask.at[:, : self.min_tokens_to_keep].set(True)
  
      topk_next_scores = jnp.where(score_mask, topk_scores, mask_scores)
      next_scores = jax.lax.sort_key_val(
          topk_indices, topk_next_scores
      )[-1]
  
      return next_scores
  
imports: import inspect
  import jax
  import jax.lax
  import jax.numpy
  from ..utils import add_start_docstrings
  from ..utils.logging import get_logger
label_desc: Use image node_id 2 for calling the FlaxTopPLogitsWarper obj's underlying member method code with example usage: obj.__call__(input_ids, scores, cur_len) and returns: next_scores