code: def test_forced_eos_token_logits_processor(self):
      vocab_size = 20
      batch_size = 4
      eos_token_id = 0
      max_length = 5
  
      logits_processor = FlaxForcedEOSTokenLogitsProcessor(
          max_length=max_length, eos_token_id=eos_token_id
      )
  
      # check that all scores are -inf except the eos_token_id when max_length is reached
      input_ids = ids_tensor((batch_size, 4), vocab_size=20)
      cur_len = 4
      scores = self._get_uniform_logits(batch_size, vocab_size)
      scores = logits_processor(input_ids, scores, cur_len=cur_len)
      self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]).all())
      self.assertListEqual(
          scores[:, eos_token_id].tolist(), 4 * [0]
      )  # score for eos_token_id should be zero
  
      # check that eos_token_id is not forced if max_length is not reached
      cur_len = 3
      scores = self._get_uniform_logits(batch_size, vocab_size)
      scores = logits_processor(input_ids, scores, cur_len=cur_len)
      self.assertFalse(jnp.isinf(scores).any())
  
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 7 for calling the LogitsProcessorTest obj's underlying member method code with example usage: obj.test_forced_eos_token_logits_processor() without return types