code: def test_top_k_dist_warper(self):
      input_ids = None
      vocab_size = 10
      batch_size = 2
  
      # create ramp distribution
      ramp_logits = np.broadcast_to(
          np.arange(vocab_size)[None, :], (batch_size, vocab_size)
      ).copy()
      ramp_logits[1:, : vocab_size // 2] = (
          ramp_logits[1:, : vocab_size // 2] + vocab_size
      )
  
      top_k_warp = FlaxTopKLogitsWarper(3)
  
      scores = top_k_warp(input_ids, ramp_logits, cur_len=None)
  
      # check that correct tokens are filtered
      self.assertListEqual(
          jnp.isinf(scores[0]).tolist(), 7 * [True] + 3 * [False]
      )
      self.assertListEqual(
          jnp.isinf(scores[1]).tolist(),
          2 * [True] + 3 * [False] + 5 * [True],
      )
  
      # check special case
      length = 5
      top_k_warp_safety_check = FlaxTopKLogitsWarper(
          top_k=1, filter_value=0.0, min_tokens_to_keep=3
      )
  
      ramp_logits = np.broadcast_to(
          np.arange(length)[None, :], (batch_size, length)
      ).copy()
      scores = top_k_warp_safety_check(
          input_ids, ramp_logits, cur_len=None
      )
  
      # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
      self.assertListEqual(
          (scores == 0.0).sum(axis=-1).tolist(), [2, 2]
      )
  
file_path: transformers/tests/generation/test_flax_logits_process.py
imports: import unittest
  import numpy
  from transformers import is_flax_available
  from transformers.testing_utils import require_flax
  from ..test_modeling_flax_common import ids_tensor
label_desc: Use image node_id 3 for calling the LogitsProcessorTest obj's underlying member method code with example usage: obj.test_top_k_dist_warper() without return types