



import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers.testing_utils import require_vision
from transformers.utils import is_vision_available


if is_vision_available():
				from PIL import Image

				from transformers import AutoProcessor, BertTokenizer, BlipImageProcessor, BlipProcessor, PreTrainedTokenizerFast

@require_vision
class  snake_case__(unittest.TestCase     ):

					"""simple docstring"""






					def 					snake_case	(				self   :      Union[str, Any]	):
									lowercase__	:       int						=	tempfile.mkdtemp()

									lowercase__	:       Union[str, Any]						=	BlipImageProcessor()
									lowercase__	:       List[Any]						=	BertTokenizer.from_pretrained("hf-internal-testing/tiny-random-BertModel"	)

									lowercase__	:       Any						=	BlipProcessor(snake_case__ ,  snake_case__	)

									processor.save_pretrained(self.tmpdirname	)







					def 					snake_case	(				self   :      int ,  **SCREAMING_SNAKE_CASE   :      Union[str, Any]	):
									return AutoProcessor.from_pretrained(self.tmpdirname ,  **snake_case__	).tokenizer







					def 					snake_case	(				self   :      str ,  **SCREAMING_SNAKE_CASE   :      int	):
									return AutoProcessor.from_pretrained(self.tmpdirname ,  **snake_case__	).image_processor







					def 					snake_case	(				self   :      List[Any]	):
									shutil.rmtree(self.tmpdirname	)







					def 					snake_case	(				self   :      Union[str, Any]	):
									lowercase__	:       Tuple						=	[np.random.randint(255 ,  size=(3, 30, 400) ,  dtype=np.uinta	)]

									lowercase__	:       Union[str, Any]						=	[Image.fromarray(np.moveaxis(snake_case__ ,  0 ,  -1	)	) for x in image_inputs]

									return image_inputs







					def 					snake_case	(				self   :      str	):
									lowercase__	:       str						=	BlipProcessor(tokenizer=self.get_tokenizer() ,  image_processor=self.get_image_processor()	)
									processor.save_pretrained(self.tmpdirname	)

									lowercase__	:       Union[str, Any]						=	self.get_tokenizer(bos_token="(BOS)" ,  eos_token="(EOS)"	)
									lowercase__	:       Union[str, Any]						=	self.get_image_processor(do_normalize=snake_case__ ,  padding_value=1.0	)

									lowercase__	:       Any						=	BlipProcessor.from_pretrained(
									    self.tmpdirname ,  bos_token="(BOS)" ,  eos_token="(EOS)" ,  do_normalize=snake_case__ ,  padding_value=1.0	)

									self.assertEqual(processor.tokenizer.get_vocab() ,  tokenizer_add_kwargs.get_vocab()	)
									self.assertIsInstance(processor.tokenizer ,  snake_case__	)

									self.assertEqual(processor.image_processor.to_json_string() ,  image_processor_add_kwargs.to_json_string()	)
									self.assertIsInstance(processor.image_processor ,  snake_case__	)







					def 					snake_case	(				self   :      Optional[int]	):
									lowercase__	:       int						=	self.get_image_processor()
									lowercase__	:       int						=	self.get_tokenizer()

									lowercase__	:       List[Any]						=	BlipProcessor(tokenizer=snake_case__ ,  image_processor=snake_case__	)

									lowercase__	:       Tuple						=	self.prepare_image_inputs()

									lowercase__	:       Optional[int]						=	image_processor(snake_case__ ,  return_tensors="np"	)
									lowercase__	:       Dict						=	processor(images=snake_case__ ,  return_tensors="np"	)

									for key in input_feat_extract.keys():
													self.assertAlmostEqual(input_feat_extract[key].sum() ,  input_processor[key].sum() ,  delta=1E-2	)







					def 					snake_case	(				self   :      Dict	):
									lowercase__	:       Union[str, Any]						=	self.get_image_processor()
									lowercase__	:       List[Any]						=	self.get_tokenizer()

									lowercase__	:       Any						=	BlipProcessor(tokenizer=snake_case__ ,  image_processor=snake_case__	)

									lowercase__	:       Optional[Any]						=	'lower newer'

									lowercase__	:       str						=	processor(text=snake_case__	)

									lowercase__	:       List[Any]						=	tokenizer(snake_case__ ,  return_token_type_ids=snake_case__	)

									for key in encoded_tok.keys():
													self.assertListEqual(encoded_tok[key] ,  encoded_processor[key]	)







					def 					snake_case	(				self   :      List[Any]	):
									lowercase__	:       Optional[int]						=	self.get_image_processor()
									lowercase__	:       Optional[int]						=	self.get_tokenizer()

									lowercase__	:       Optional[int]						=	BlipProcessor(tokenizer=snake_case__ ,  image_processor=snake_case__	)

									lowercase__	:       Tuple						=	'lower newer'
									lowercase__	:       Tuple						=	self.prepare_image_inputs()

									lowercase__	:       List[Any]						=	processor(text=snake_case__ ,  images=snake_case__	)

									self.assertListEqual(list(inputs.keys()	) ,  ["pixel_values", "input_ids", "attention_mask"]	)

									# test if it raises when no input is passed
									with pytest.raises(snake_case__	):
													processor()







					def 					snake_case	(				self   :      Optional[int]	):
									lowercase__	:       Any						=	self.get_image_processor()
									lowercase__	:       str						=	self.get_tokenizer()

									lowercase__	:       Optional[Any]						=	BlipProcessor(tokenizer=snake_case__ ,  image_processor=snake_case__	)

									lowercase__	:       Optional[int]						=	[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

									lowercase__	:       str						=	processor.batch_decode(snake_case__	)
									lowercase__	:       Any						=	tokenizer.batch_decode(snake_case__	)

									self.assertListEqual(snake_case__ ,  snake_case__	)







					def 					snake_case	(				self   :      Union[str, Any]	):
									lowercase__	:       List[str]						=	self.get_image_processor()
									lowercase__	:       List[Any]						=	self.get_tokenizer()

									lowercase__	:       Dict						=	BlipProcessor(tokenizer=snake_case__ ,  image_processor=snake_case__	)

									lowercase__	:       Optional[int]						=	'lower newer'
									lowercase__	:       List[Any]						=	self.prepare_image_inputs()

									lowercase__	:       Optional[int]						=	processor(text=snake_case__ ,  images=snake_case__	)

									# For now the processor supports only ['pixel_values', 'input_ids', 'attention_mask']
									self.assertListEqual(list(inputs.keys()	) ,  ["pixel_values", "input_ids", "attention_mask"]	)










import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)







@require_flax
class 				lowerCAmelCase_       (					unittest.TestCase							):






				def 		snake_case	(			self      ,snake_case__      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      jnp.ones((batch_size, length) ) / length
							return scores






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      None
							SCREAMING_SNAKE_CASE_   : int					  =      20

							SCREAMING_SNAKE_CASE_   : int					  =      self._get_uniform_logits(batch_size=2      ,length=snake_case__ )

							# tweak scores to not be uniform anymore
							SCREAMING_SNAKE_CASE_   : Any					  =      scores.at[1, 5].set((1 / length) + 0.1 )  # peak, 1st batch
							SCREAMING_SNAKE_CASE_   : Dict					  =      scores.at[1, 10].set((1 / length) - 0.4 )  # valley, 1st batch

							# compute softmax
							SCREAMING_SNAKE_CASE_   : int					  =      jax.nn.softmax(snake_case__      ,axis=-1 )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTemperatureLogitsWarper(temperature=1.3 )

							SCREAMING_SNAKE_CASE_   : Any					  =      jax.nn.softmax(temp_dist_warper_sharper(snake_case__      ,scores.copy()      ,cur_len=snake_case__ )      ,axis=-1 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      jax.nn.softmax(temp_dist_warper_smoother(snake_case__      ,scores.copy()      ,cur_len=snake_case__ )      ,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() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Any					  =      None
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      10
							SCREAMING_SNAKE_CASE_   : int					  =      2

							# create ramp distribution
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, vocab_size) ).copy()
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      ramp_logits[1:, : vocab_size // 2] + vocab_size

							SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxTopKLogitsWarper(3 )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      top_k_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# 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
							SCREAMING_SNAKE_CASE_   : int					  =      5
							SCREAMING_SNAKE_CASE_   : Any					  =      FlaxTopKLogitsWarper(top_k=1      ,filter_value=0.0      ,min_tokens_to_keep=3 )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, length) ).copy()
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      top_k_warp_safety_check(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
							self.assertListEqual((scores == 0.0).sum(axis=-1 ).tolist()      ,[2, 2] )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      None
							SCREAMING_SNAKE_CASE_   : str					  =      10
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      2

							# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
							SCREAMING_SNAKE_CASE_   : Any					  =      np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]] ) )

							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTopPLogitsWarper(0.8 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      np.exp(top_p_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ ) )

							# dist should be filtered to keep min num values so that sum is >= top_p
							# exp (-inf) => 0
							SCREAMING_SNAKE_CASE_   : str					  =      np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]] )
							self.assertTrue(np.allclose(snake_case__      ,snake_case__      ,atol=1E-3 ) )

							# check edge cases with negative and extreme logits
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      np.broadcast_to(np.arange(snake_case__ )[None, :]      ,(batch_size, vocab_size) ).copy() - (
							    vocab_size // 2
							)

							# make ramp_logits more extreme
							SCREAMING_SNAKE_CASE_   : Any					  =      ramp_logits[1] * 100.0

							# make sure at least 2 tokens are kept
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      FlaxTopPLogitsWarper(0.9      ,min_tokens_to_keep=2      ,filter_value=0.0 )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      top_p_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# 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] )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      20
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      4
							SCREAMING_SNAKE_CASE_   : str					  =      0

							SCREAMING_SNAKE_CASE_   : str					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )

							# check that min length is applied at length 5
							SCREAMING_SNAKE_CASE_   : Dict					  =      ids_tensor((batch_size, 20)      ,vocab_size=20 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      5
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      min_dist_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()      ,4 * [-float('inf' )] )

							# check that min length is not applied anymore at length 15
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      15
							SCREAMING_SNAKE_CASE_   : Any					  =      min_dist_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertFalse(jnp.isinf(snake_case__ ).any() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      20
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      0

							SCREAMING_SNAKE_CASE_   : int					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )

							# check that all scores are -inf except the bos_token_id score
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ids_tensor((batch_size, 1)      ,vocab_size=20 )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      1
							SCREAMING_SNAKE_CASE_   : str					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      logits_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :] ).all() )
							self.assertListEqual(scores[:, bos_token_id].tolist()      ,4 * [0] )  # score for bos_token_id shold be zero

							# check that bos_token_id is not forced if current length is greater than 1
							SCREAMING_SNAKE_CASE_   : Dict					  =      3
							SCREAMING_SNAKE_CASE_   : Dict					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      logits_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertFalse(jnp.isinf(snake_case__ ).any() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      20
							SCREAMING_SNAKE_CASE_   : Any					  =      4
							SCREAMING_SNAKE_CASE_   : int					  =      0
							SCREAMING_SNAKE_CASE_   : Tuple					  =      5

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,eos_token_id=snake_case__ )

							# check that all scores are -inf except the eos_token_id when max_length is reached
							SCREAMING_SNAKE_CASE_   : Any					  =      ids_tensor((batch_size, 4)      ,vocab_size=20 )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      logits_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							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
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      3
							SCREAMING_SNAKE_CASE_   : int					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      logits_processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							self.assertFalse(jnp.isinf(snake_case__ ).any() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      4
							SCREAMING_SNAKE_CASE_   : Dict					  =      10
							SCREAMING_SNAKE_CASE_   : int					  =      15
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      2
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      1
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      15

							# dummy input_ids and scores
							SCREAMING_SNAKE_CASE_   : Any					  =      ids_tensor((batch_size, sequence_length)      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      input_ids.copy()

							SCREAMING_SNAKE_CASE_   : str					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      scores.copy()

							# instantiate all dist processors
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTopKLogitsWarper(3 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      FlaxTopPLogitsWarper(0.8 )

							# instantiate all logits processors
							SCREAMING_SNAKE_CASE_   : Any					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,eos_token_id=snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      10

							# no processor list
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      temp_dist_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      top_k_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      top_p_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : str					  =      min_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      bos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      eos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# with processor list
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
							SCREAMING_SNAKE_CASE_   : Any					  =      processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )

							# scores should be equal
							self.assertTrue(jnp.allclose(snake_case__      ,snake_case__      ,atol=1E-3 ) )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()      ,input_ids_comp.tolist() )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      4
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      10
							SCREAMING_SNAKE_CASE_   : Dict					  =      15
							SCREAMING_SNAKE_CASE_   : Dict					  =      2
							SCREAMING_SNAKE_CASE_   : Tuple					  =      1
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      15

							# dummy input_ids and scores
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ids_tensor((batch_size, sequence_length)      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      input_ids.copy()

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self._get_uniform_logits(snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      scores.copy()

							# instantiate all dist processors
							SCREAMING_SNAKE_CASE_   : str					  =      FlaxTemperatureLogitsWarper(temperature=0.5 )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxTopKLogitsWarper(3 )
							SCREAMING_SNAKE_CASE_   : int					  =      FlaxTopPLogitsWarper(0.8 )

							# instantiate all logits processors
							SCREAMING_SNAKE_CASE_   : Tuple					  =      FlaxMinLengthLogitsProcessor(min_length=10      ,eos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
							SCREAMING_SNAKE_CASE_   : str					  =      FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,eos_token_id=snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      10

							# no processor list
							def run_no_processor_list(snake_case__      ,snake_case__      ,snake_case__ ):
										SCREAMING_SNAKE_CASE_   : int					  =      temp_dist_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										SCREAMING_SNAKE_CASE_   : int					  =      top_k_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      top_p_warp(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      min_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      bos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      eos_dist_proc(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										return scores

							# with processor list
							def run_processor_list(snake_case__      ,snake_case__      ,snake_case__ ):
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      FlaxLogitsProcessorList(
										    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
										SCREAMING_SNAKE_CASE_   : List[str]					  =      processor(snake_case__      ,snake_case__      ,cur_len=snake_case__ )
										return scores

							SCREAMING_SNAKE_CASE_   : Tuple					  =      jax.jit(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      jax.jit(snake_case__ )

							SCREAMING_SNAKE_CASE_   : int					  =      jitted_run_no_processor_list(snake_case__      ,snake_case__      ,snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[str]					  =      jitted_run_processor_list(snake_case__      ,snake_case__      ,snake_case__ )

							# scores should be equal
							self.assertTrue(jnp.allclose(snake_case__      ,snake_case__      ,atol=1E-3 ) )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()      ,input_ids_comp.tolist() )



