code: def test_top_p_dist_warper(self):
      input_ids = None
      vocab_size = 10
      batch_size = 2
  
      # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
      dist = np.log(
          np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]])
      )
  
      top_p_warp = FlaxTopPLogitsWarper(0.8)
      filtered_dist = np.exp(top_p_warp(input_ids, dist, cur_len=None))
  
      # dist should be filtered to keep min num values so that sum is >= top_p
      # exp (-inf) => 0
      EXPECTED_FILTERED_DIST = np.array(
          [[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]
      )
      self.assertTrue(
          np.allclose(filtered_dist, EXPECTED_FILTERED_DIST, atol=1e-3)
      )
  
      # check edge cases with negative and extreme logits
      ramp_logits = np.broadcast_to(
          np.arange(vocab_size)[None, :], (batch_size, vocab_size)
      ).copy() - (vocab_size // 2)
  
      # make ramp_logits more extreme
      ramp_logits[1] = ramp_logits[1] * 100.0
  
      # make sure at least 2 tokens are kept
      top_p_warp = FlaxTopPLogitsWarper(
          0.9, min_tokens_to_keep=2, filter_value=0.0
      )
      filtered_dist = top_p_warp(input_ids, ramp_logits, cur_len=None)
  
      # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
      self.assertListEqual(
          (filtered_dist != 0.0).sum(axis=-1).tolist(), [3, 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 4 for calling the LogitsProcessorTest obj's underlying member method code with example usage: obj.test_top_p_dist_warper() without return types