class_docstr: [`FlaxLogitsWarper`] that performs top-k, i.e. restricting to the k highest probability elements.
  
  Args:
      top_k (`int`):
          The number of highest probability vocabulary tokens to keep for top-k-filtering.
      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:
      batch_size, vocab_size = scores.shape
      next_scores_flat = jnp.full(
          batch_size * vocab_size, self.filter_value
      )
  
      topk = min(self.top_k, scores.shape[-1])  # Safety check
      topk_scores, topk_indices = lax.top_k(scores, topk)
      shift = jnp.broadcast_to(
          (jnp.arange(batch_size) * vocab_size)[:, None],
          (batch_size, topk),
      ).flatten()
      topk_scores_flat = topk_scores.flatten()
      topk_indices_flat = topk_indices.flatten() + shift
  
      next_scores_flat = next_scores_flat.at[topk_indices_flat].set(
          topk_scores_flat
      )
      next_scores = next_scores_flat.reshape(batch_size, vocab_size)
      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 FlaxTopKLogitsWarper obj's underlying member method code with example usage: obj.__call__(input_ids, scores, cur_len) and returns: next_scores