



'''simple docstring'''




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    A_	(unittest.TestCase							):
							"""simple docstring"""




							def 							_A	(		self			:int       ,							lowerCAmelCase__			:List[Any]       ,							lowerCAmelCase__			:Optional[int]				)				->  List[str]:


										'''simple docstring'''




										snake_case_				:					Optional[Any]   =       jnp.ones((batch_size, length)				) / length
										return scores

							def 							_A	(		self			:Any				)				->  List[Any]:


										'''simple docstring'''




										snake_case_				:					Optional[Any]   =       None
										snake_case_				:					Optional[int]   =       20

										snake_case_				:					Tuple   =       self._get_uniform_logits(batch_size=2       ,							length=lowerCAmelCase__				)

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

										# compute softmax
										snake_case_				:					Any   =       jax.nn.softmax(lowerCAmelCase__       ,							axis=-1				)

										snake_case_				:					Tuple   =       FlaxTemperatureLogitsWarper(temperature=0.5				)
										snake_case_				:					Optional[int]   =       FlaxTemperatureLogitsWarper(temperature=1.3				)

										snake_case_				:					Union[str, Any]   =       jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase__       ,							scores.copy()       ,							cur_len=lowerCAmelCase__				)       ,							axis=-1				)
										snake_case_				:					Tuple   =       jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase__       ,							scores.copy()       ,							cur_len=lowerCAmelCase__				)       ,							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 							_A	(		self			:Optional[int]				)				->  Any:


										'''simple docstring'''




										snake_case_				:					Tuple   =       None
										snake_case_				:					Any   =       10
										snake_case_				:					int   =       2

										# create ramp distribution
										snake_case_				:					Dict   =       np.broadcast_to(np.arange(lowerCAmelCase__				)[None, :]       ,							(batch_size, vocab_size)				).copy()
										snake_case_				:					Dict   =       ramp_logits[1:, : vocab_size // 2] + vocab_size

										snake_case_				:					Dict   =       FlaxTopKLogitsWarper(3				)

										snake_case_				:					int   =       top_k_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)

										# 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
										snake_case_				:					str   =       5
										snake_case_				:					Optional[int]   =       FlaxTopKLogitsWarper(top_k=1       ,							filter_value=0.0       ,							min_tokens_to_keep=3				)

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

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

							def 							_A	(		self			:List[str]				)				->  List[str]:


										'''simple docstring'''




										snake_case_				:					Optional[Any]   =       None
										snake_case_				:					List[str]   =       10
										snake_case_				:					Tuple   =       2

										# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
										snake_case_				:					Optional[Any]   =       np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]				)				)

										snake_case_				:					int   =       FlaxTopPLogitsWarper(0.8				)
										snake_case_				:					Dict   =       np.exp(top_p_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)				)

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

										# check edge cases with negative and extreme logits
										snake_case_				:					Dict   =       np.broadcast_to(np.arange(lowerCAmelCase__				)[None, :]       ,							(batch_size, vocab_size)				).copy() - (
										    vocab_size // 2
										)

										# make ramp_logits more extreme
										snake_case_				:					Dict   =       ramp_logits[1] * 1_0_0.0

										# make sure at least 2 tokens are kept
										snake_case_				:					int   =       FlaxTopPLogitsWarper(0.9       ,							min_tokens_to_keep=2       ,							filter_value=0.0				)
										snake_case_				:					Optional[Any]   =       top_p_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)

										# 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 							_A	(		self			:int				)				->  List[Any]:


										'''simple docstring'''




										snake_case_				:					Union[str, Any]   =       20
										snake_case_				:					Optional[int]   =       4
										snake_case_				:					Dict   =       0

										snake_case_				:					List[str]   =       FlaxMinLengthLogitsProcessor(min_length=10       ,							eos_token_id=lowerCAmelCase__				)

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

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

							def 							_A	(		self			:Optional[int]				)				->  List[Any]:


										'''simple docstring'''




										snake_case_				:					Union[str, Any]   =       20
										snake_case_				:					Tuple   =       4
										snake_case_				:					Dict   =       0

										snake_case_				:					List[str]   =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__				)

										# check that all scores are -inf except the bos_token_id score
										snake_case_				:					List[str]   =       ids_tensor((batch_size, 1)       ,							vocab_size=20				)
										snake_case_				:					Optional[int]   =       1
										snake_case_				:					List[Any]   =       self._get_uniform_logits(lowerCAmelCase__       ,							lowerCAmelCase__				)
										snake_case_				:					Tuple   =       logits_processor(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										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
										snake_case_				:					Tuple   =       3
										snake_case_				:					List[Any]   =       self._get_uniform_logits(lowerCAmelCase__       ,							lowerCAmelCase__				)
										snake_case_				:					str   =       logits_processor(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										self.assertFalse(jnp.isinf(lowerCAmelCase__				).any()				)

							def 							_A	(		self			:str				)				->  Union[str, Any]:


										'''simple docstring'''




										snake_case_				:					Dict   =       20
										snake_case_				:					Union[str, Any]   =       4
										snake_case_				:					int   =       0
										snake_case_				:					Union[str, Any]   =       5

										snake_case_				:					int   =       FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       ,							eos_token_id=lowerCAmelCase__				)

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

							def 							_A	(		self			:Optional[int]				)				->  int:


										'''simple docstring'''




										snake_case_				:					Optional[Any]   =       4
										snake_case_				:					Optional[Any]   =       10
										snake_case_				:					Union[str, Any]   =       15
										snake_case_				:					Union[str, Any]   =       2
										snake_case_				:					Any   =       1
										snake_case_				:					Any   =       15

										# dummy input_ids and scores
										snake_case_				:					str   =       ids_tensor((batch_size, sequence_length)       ,							lowerCAmelCase__				)
										snake_case_				:					int   =       input_ids.copy()

										snake_case_				:					List[Any]   =       self._get_uniform_logits(lowerCAmelCase__       ,							lowerCAmelCase__				)
										snake_case_				:					List[Any]   =       scores.copy()

										# instantiate all dist processors
										snake_case_				:					str   =       FlaxTemperatureLogitsWarper(temperature=0.5				)
										snake_case_				:					Optional[Any]   =       FlaxTopKLogitsWarper(3				)
										snake_case_				:					str   =       FlaxTopPLogitsWarper(0.8				)

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

										snake_case_				:					Optional[Any]   =       10

										# no processor list
										snake_case_				:					List[Any]   =       temp_dist_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										snake_case_				:					List[Any]   =       top_k_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										snake_case_				:					List[Any]   =       top_p_warp(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										snake_case_				:					Optional[Any]   =       min_dist_proc(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										snake_case_				:					int   =       bos_dist_proc(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
										snake_case_				:					Optional[int]   =       eos_dist_proc(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)

										# with processor list
										snake_case_				:					List[str]   =       FlaxLogitsProcessorList(
										    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]				)
										snake_case_				:					Optional[int]   =       processor(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)

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

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






							def 							_A	(		self			:Optional[Any]				)				->  Optional[Any]:


										'''simple docstring'''




										snake_case_				:					Optional[Any]   =       4
										snake_case_				:					Any   =       10
										snake_case_				:					Optional[Any]   =       15
										snake_case_				:					int   =       2
										snake_case_				:					List[Any]   =       1
										snake_case_				:					Dict   =       15

										# dummy input_ids and scores
										snake_case_				:					List[str]   =       ids_tensor((batch_size, sequence_length)       ,							lowerCAmelCase__				)
										snake_case_				:					List[str]   =       input_ids.copy()

										snake_case_				:					int   =       self._get_uniform_logits(lowerCAmelCase__       ,							lowerCAmelCase__				)
										snake_case_				:					List[str]   =       scores.copy()

										# instantiate all dist processors
										snake_case_				:					Union[str, Any]   =       FlaxTemperatureLogitsWarper(temperature=0.5				)
										snake_case_				:					Union[str, Any]   =       FlaxTopKLogitsWarper(3				)
										snake_case_				:					Tuple   =       FlaxTopPLogitsWarper(0.8				)

										# instantiate all logits processors
										snake_case_				:					int   =       FlaxMinLengthLogitsProcessor(min_length=10       ,							eos_token_id=lowerCAmelCase__				)
										snake_case_				:					Optional[int]   =       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__				)
										snake_case_				:					Any   =       FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       ,							eos_token_id=lowerCAmelCase__				)

										snake_case_				:					List[Any]   =       10

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

										# with processor list
										def run_processor_list(lowerCAmelCase__			:str       ,							lowerCAmelCase__			:int       ,							lowerCAmelCase__			:Tuple				):
													snake_case_				:					Union[str, Any]   =       FlaxLogitsProcessorList(
													    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]				)
													snake_case_				:					List[str]   =       processor(lowerCAmelCase__       ,							lowerCAmelCase__       ,							cur_len=lowerCAmelCase__				)
													return scores

										snake_case_				:					Tuple   =       jax.jit(lowerCAmelCase__				)
										snake_case_				:					str   =       jax.jit(lowerCAmelCase__				)

										snake_case_				:					Any   =       jitted_run_no_processor_list(lowerCAmelCase__       ,							lowerCAmelCase__       ,							lowerCAmelCase__				)
										snake_case_				:					Dict   =       jitted_run_processor_list(lowerCAmelCase__       ,							lowerCAmelCase__       ,							lowerCAmelCase__				)

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

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












'''simple docstring'''




def 	__UpperCAmelCase       ( __magic_name__						)->    list[int]:



			"""simple docstring"""



			if length <= 0 or not isinstance(__magic_name__     ,__magic_name__						):
						raise ValueError("Length must be a positive integer."						)
			return [n * (2 * n - 1) for n in range(__magic_name__						)]


if __name__ == "__main__":
		print(hexagonal_numbers(length=5))
		print(hexagonal_numbers(length=10))




