code: def test_temperature_dist_warper(self):
      input_ids = None
      length = 20
  
      scores = self._get_uniform_logits(batch_size=2, length=length)
  
      # tweak scores to not be uniform anymore
      scores = scores.at[1, 5].set(
          (1 / length) + 0.1
      )  # peak, 1st batch
      scores = scores.at[1, 10].set(
          (1 / length) - 0.4
      )  # valley, 1st batch
  
      # compute softmax
      probs = jax.nn.softmax(scores, axis=-1)
  
      temp_dist_warper_sharper = FlaxTemperatureLogitsWarper(
          temperature=0.5
      )
      temp_dist_warper_smoother = FlaxTemperatureLogitsWarper(
          temperature=1.3
      )
  
      warped_prob_sharp = jax.nn.softmax(
          temp_dist_warper_sharper(
              input_ids, scores.copy(), cur_len=None
          ),
          axis=-1,
      )
      warped_prob_smooth = jax.nn.softmax(
          temp_dist_warper_smoother(
              input_ids, scores.copy(), cur_len=None
          ),
          axis=-1,
      )
  
      # uniform distribution stays uniform
      self.assertTrue(
          jnp.allclose(probs[0, :], warped_prob_sharp[0, :], atol=1e-3)
      )
      self.assertTrue(
          jnp.allclose(probs[0, :], warped_prob_smooth[0, :], atol=1e-3)
      )
  
      # sharp peaks get higher, valleys get lower
      self.assertLess(probs[1, :].max(), warped_prob_sharp[1, :].max())
      self.assertGreater(
          probs[1, :].min(), warped_prob_sharp[1, :].min()
      )
  
      # smooth peaks get lower, valleys get higher
      self.assertGreater(
          probs[1, :].max(), warped_prob_smooth[1, :].max()
      )
      self.assertLess(probs[1, :].min(), warped_prob_smooth[1, :].min())
  
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 2 for calling the LogitsProcessorTest obj's underlying member method code with example usage: obj.test_temperature_dist_warper() without return types